[InstCombine] enhance (select X, C1, C2 --> ext X) to handle vectors

By replacing dyn_cast of ConstantInt with m_Zero/m_One/m_AllOnes, we
allow these transforms for splat vectors.

Differential Revision: http://reviews.llvm.org/D21899

llvm-svn: 274696

2 files changed, 46 insertions, 27 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index d89691e5b92..db4ef6ec22d 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -954,32 +954,38 @@ Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
       return BinaryOperator::CreateOr(TrueVal, FalseVal);
   }
 
-  // Selecting between two integer constants?
-  if (ConstantInt *TrueValC = dyn_cast<ConstantInt>(TrueVal))
-    if (ConstantInt *FalseValC = dyn_cast<ConstantInt>(FalseVal)) {
-      // select C, 1, 0 -> zext C to int
-      if (FalseValC->isZero() && TrueValC->getValue() == 1)
-        return new ZExtInst(CondVal, SI.getType());
-
-      // select C, -1, 0 -> sext C to int
-      if (FalseValC->isZero() && TrueValC->isAllOnesValue())
-        return new SExtInst(CondVal, SI.getType());
-
-      // select C, 0, 1 -> zext !C to int
-      if (TrueValC->isZero() && FalseValC->getValue() == 1) {
-        Value *NotCond = Builder->CreateNot(CondVal, "not."+CondVal->getName());
-        return new ZExtInst(NotCond, SI.getType());
-      }
+  // Selecting between two integer or vector splat integer constants?
+  //
+  // Note that we don't handle a scalar select of vectors:
+  // select i1 %c, <2 x i8> <1, 1>, <2 x i8> <0, 0>
+  // because that may need 3 instructions to splat the condition value:
+  // extend, insertelement, shufflevector.
+  if (CondVal->getType()->isVectorTy() == SI.getType()->isVectorTy()) {
+    // select C, 1, 0 -> zext C to int
+    if (match(TrueVal, m_One()) && match(FalseVal, m_Zero()))
+      return new ZExtInst(CondVal, SI.getType());
+
+    // select C, -1, 0 -> sext C to int
+    if (match(TrueVal, m_AllOnes()) && match(FalseVal, m_Zero()))
+      return new SExtInst(CondVal, SI.getType());
+
+    // select C, 0, 1 -> zext !C to int
+    if (match(TrueVal, m_Zero()) && match(FalseVal, m_One())) {
+      Value *NotCond = Builder->CreateNot(CondVal, "not." + CondVal->getName());
+      return new ZExtInst(NotCond, SI.getType());
+    }
 
-      // select C, 0, -1 -> sext !C to int
-      if (TrueValC->isZero() && FalseValC->isAllOnesValue()) {
-        Value *NotCond = Builder->CreateNot(CondVal, "not."+CondVal->getName());
-        return new SExtInst(NotCond, SI.getType());
-      }
+    // select C, 0, -1 -> sext !C to int
+    if (match(TrueVal, m_Zero()) && match(FalseVal, m_AllOnes())) {
+      Value *NotCond = Builder->CreateNot(CondVal, "not." + CondVal->getName());
+      return new SExtInst(NotCond, SI.getType());
+    }
+  }
 
+  if (ConstantInt *TrueValC = dyn_cast<ConstantInt>(TrueVal))
+    if (ConstantInt *FalseValC = dyn_cast<ConstantInt>(FalseVal))
       if (Value *V = foldSelectICmpAnd(SI, TrueValC, FalseValC, Builder))
         return replaceInstUsesWith(SI, V);
-    }
 
   // See if we are selecting two values based on a comparison of the two values.
   if (FCmpInst *FCI = dyn_cast<FCmpInst>(CondVal)) {
diff --git a/llvm/test/Transforms/InstCombine/apint-select.ll b/llvm/test/Transforms/InstCombine/apint-select.ll
index 049da268e5d..2aae232f1b0 100644
--- a/llvm/test/Transforms/InstCombine/apint-select.ll
+++ b/llvm/test/Transforms/InstCombine/apint-select.ll
@@ -41,11 +41,11 @@ define i999 @not_sext(i1 %C) {
   ret i999 %V
 }
 
-; FIXME: Vector selects of vector splat constants match APInt too.
+; Vector selects of vector splat constants match APInt too.
 
 define <2 x i41> @zext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @zext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i41> <i41 1, i41 1>, <2 x i41> zeroinitializer
+; CHECK-NEXT:    [[V:%.*]] = zext <2 x i1> %C to <2 x i41>
 ; CHECK-NEXT:    ret <2 x i41> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i41> <i41 1, i41 1>, <2 x i41> <i41 0, i41 0>
@@ -54,7 +54,7 @@ define <2 x i41> @zext_vec(<2 x i1> %C) {
 
 define <2 x i32> @sext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @sext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i32> <i32 -1, i32 -1>, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[V:%.*]] = sext <2 x i1> %C to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i32> <i32 -1, i32 -1>, <2 x i32> <i32 0, i32 0>
@@ -63,7 +63,8 @@ define <2 x i32> @sext_vec(<2 x i1> %C) {
 
 define <2 x i999> @not_zext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @not_zext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i999> zeroinitializer, <2 x i999> <i999 1, i999 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i1> %C to <2 x i999>
+; CHECK-NEXT:    [[V:%.*]] = xor <2 x i999> [[TMP1]], <i999 1, i999 1>
 ; CHECK-NEXT:    ret <2 x i999> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i999> <i999 0, i999 0>, <2 x i999> <i999 1, i999 1>
@@ -72,13 +73,25 @@ define <2 x i999> @not_zext_vec(<2 x i1> %C) {
 
 define <2 x i64> @not_sext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @not_sext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i64> zeroinitializer, <2 x i64> <i64 -1, i64 -1>
+; CHECK-NEXT:    [[NOT_C:%.*]] = xor <2 x i1> %C, <i1 true, i1 true>
+; CHECK-NEXT:    [[V:%.*]] = sext <2 x i1> [[NOT_C]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i64> <i64 0, i64 0>, <2 x i64> <i64 -1, i64 -1>
   ret <2 x i64> %V
 }
 
+; But don't touch this - we would need 3 instructions to extend and splat the scalar select condition.
+
+define <2 x i32> @scalar_select_of_vectors(i1 %c) {
+; CHECK-LABEL: @scalar_select_of_vectors(
+; CHECK-NEXT:    [[V:%.*]] = select i1 %c, <2 x i32> <i32 1, i32 1>, <2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <2 x i32> [[V]]
+;
+  %V = select i1 %c, <2 x i32> <i32 1, i32 1>, <2 x i32> zeroinitializer
+  ret <2 x i32> %V
+}
+
 ;; (x <s 0) ? -1 : 0 -> ashr x, 31
 
 define i41 @test3(i41 %X) {