[InstCombine] fold select X, (ext X), C

If we're going to canonicalize IR towards select of constants, try harder to create those.
Also, don't lose the metadata.

This is actually 4 related transforms in one patch:
      // select X, (sext X), C --> select X, -1, C
      // select X, (zext X), C --> select X,  1, C
      // select X, C, (sext X) --> select X, C, 0
      // select X, C, (zext X) --> select X, C, 0

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

llvm-svn: 283575

2 files changed, 29 insertions, 17 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index b8213426a63..042b3731a90 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -939,11 +939,11 @@ Instruction *InstCombiner::foldSelectExtConst(SelectInst &Sel) {
 
   // If the constant is the same after truncation to the smaller type and
   // extension to the original type, we can narrow the select.
+  Value *Cond = Sel.getCondition();
   Type *SelType = Sel.getType();
   Constant *TruncC = ConstantExpr::getTrunc(C, SmallType);
   Constant *ExtC = ConstantExpr::getCast(ExtOpcode, TruncC, SelType);
   if (ExtC == C) {
-    Value *Cond = Sel.getCondition();
     Value *TruncCVal = cast<Value>(TruncC);
     if (ExtInst == Sel.getFalseValue())
       std::swap(X, TruncCVal);
@@ -954,6 +954,26 @@ Instruction *InstCombiner::foldSelectExtConst(SelectInst &Sel) {
     return CastInst::Create(Instruction::CastOps(ExtOpcode), NewSel, SelType);
   }
 
+  // If one arm of the select is the extend of the condition, replace that arm
+  // with the extension of the appropriate known bool value.
+  if (Cond == X) {
+    SelectInst *NewSel;
+    if (ExtInst == Sel.getTrueValue()) {
+      // select X, (sext X), C --> select X, -1, C
+      // select X, (zext X), C --> select X,  1, C
+      Constant *One = ConstantInt::getTrue(SmallType);
+      Constant *AllOnesOrOne = ConstantExpr::getCast(ExtOpcode, One, SelType);
+      NewSel = SelectInst::Create(Cond, AllOnesOrOne, C);
+    } else {
+      // select X, C, (sext X) --> select X, C, 0
+      // select X, C, (zext X) --> select X, C, 0
+      Constant *Zero = ConstantInt::getNullValue(SelType);
+      NewSel = SelectInst::Create(Cond, C, Zero);
+    }
+    NewSel->copyMetadata(Sel);
+    return NewSel;
+  }
+
   return nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/select-bitext.ll b/llvm/test/Transforms/InstCombine/select-bitext.ll
index 3fdba41ad29..2717e0c4560 100644
--- a/llvm/test/Transforms/InstCombine/select-bitext.ll
+++ b/llvm/test/Transforms/InstCombine/select-bitext.ll
@@ -223,8 +223,7 @@ define <2 x i32> @scalar_select_of_vectors_zext(<2 x i1> %cca, i1 %ccb) {
 
 define i32 @sext_true_val_must_be_all_ones(i1 %x) {
 ; CHECK-LABEL: @sext_true_val_must_be_all_ones(
-; CHECK-NEXT:    [[EXT:%.*]] = sext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 [[EXT]], i32 42, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 -1, i32 42, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = sext i1 %x to i32
@@ -234,8 +233,7 @@ define i32 @sext_true_val_must_be_all_ones(i1 %x) {
 
 define <2 x i32> @sext_true_val_must_be_all_ones_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @sext_true_val_must_be_all_ones_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = sext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> [[EXT]], <2 x i32> <i32 42, i32 12>, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 -1, i32 -1>, <2 x i32> <i32 42, i32 12>, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = sext <2 x i1> %x to <2 x i32>
@@ -245,8 +243,7 @@ define <2 x i32> @sext_true_val_must_be_all_ones_vec(<2 x i1> %x) {
 
 define i32 @zext_true_val_must_be_one(i1 %x) {
 ; CHECK-LABEL: @zext_true_val_must_be_one(
-; CHECK-NEXT:    [[EXT:%.*]] = zext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 [[EXT]], i32 42, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 1, i32 42, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = zext i1 %x to i32
@@ -256,8 +253,7 @@ define i32 @zext_true_val_must_be_one(i1 %x) {
 
 define <2 x i32> @zext_true_val_must_be_one_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @zext_true_val_must_be_one_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = zext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> [[EXT]], <2 x i32> <i32 42, i32 12>, !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 1, i32 1>, <2 x i32> <i32 42, i32 12>, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = zext <2 x i1> %x to <2 x i32>
@@ -267,8 +263,7 @@ define <2 x i32> @zext_true_val_must_be_one_vec(<2 x i1> %x) {
 
 define i32 @sext_false_val_must_be_zero(i1 %x) {
 ; CHECK-LABEL: @sext_false_val_must_be_zero(
-; CHECK-NEXT:    [[EXT:%.*]] = sext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 0, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = sext i1 %x to i32
@@ -278,8 +273,7 @@ define i32 @sext_false_val_must_be_zero(i1 %x) {
 
 define <2 x i32> @sext_false_val_must_be_zero_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @sext_false_val_must_be_zero_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = sext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> zeroinitializer, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = sext <2 x i1> %x to <2 x i32>
@@ -289,8 +283,7 @@ define <2 x i32> @sext_false_val_must_be_zero_vec(<2 x i1> %x) {
 
 define i32 @zext_false_val_must_be_zero(i1 %x) {
 ; CHECK-LABEL: @zext_false_val_must_be_zero(
-; CHECK-NEXT:    [[EXT:%.*]] = zext i1 %x to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 %x, i32 42, i32 0, !prof !0
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %ext = zext i1 %x to i32
@@ -300,8 +293,7 @@ define i32 @zext_false_val_must_be_zero(i1 %x) {
 
 define <2 x i32> @zext_false_val_must_be_zero_vec(<2 x i1> %x) {
 ; CHECK-LABEL: @zext_false_val_must_be_zero_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = zext <2 x i1> %x to <2 x i32>
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> [[EXT]], !prof !0
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> %x, <2 x i32> <i32 42, i32 12>, <2 x i32> zeroinitializer, !prof !0
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %ext = zext <2 x i1> %x to <2 x i32>