[InstCombine] Handle integer extension in `select` patterns using the condition as value

These patterns were previously only implemented for i1 type but can be extended for any integer type by also handling zext and sext operands.

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

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

@@ -2662,6 +2662,51 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   if (Instruction *I = canonicalizeScalarSelectOfVecs(SI, *this))
     return I;
 
+  // If the type of select is not an integer type or if the condition and
+  // the selection type are not both scalar nor both vector types, there is no
+  // point in attempting to match these patterns.
+  if (!isa<Constant>(CondVal) && SelType->isIntOrIntVectorTy() &&
+      CondVal->getType()->isVectorTy() == SelType->isVectorTy()) {
+    auto *One = ConstantInt::get(SelType, 1);
+    auto *Zero = ConstantInt::get(SelType, 0);
+    auto *AllOnes = ConstantInt::get(SelType, -1, /*isSigned*/ true);
+
+    // select a, a, b -> select a, 1, b
+    // select a, zext(a), b -> select a, 1, b
+    if (match(TrueVal, m_ZExtOrSelf(m_Specific(CondVal))))
+      return replaceOperand(SI, 1, One);
+
+    // select a, sext(a), b -> select a, -1, b
+    if (match(TrueVal, m_SExt(m_Specific(CondVal))))
+      return replaceOperand(SI, 1, AllOnes);
+
+    // select a, b, a -> select a, b, 0
+    // select a, b, zext(a) -> select a, b, 0
+    // select a, b, sext(a) -> select a, b, 0
+    if (match(FalseVal, m_ZExtOrSExtOrSelf(m_Specific(CondVal))))
+      return replaceOperand(SI, 2, Zero);
+
+    Value *NotCond;
+
+    // select a, !a, b -> select !a, b, 0
+    // select a, sext(!a), b -> select !a, b, 0
+    // select a, zext(!a), b -> select !a, b, 0
+    if (match(TrueVal, m_ZExtOrSExtOrSelf(m_CombineAnd(
+                           m_Value(NotCond), m_Not(m_Specific(CondVal))))))
+      return SelectInst::Create(NotCond, FalseVal, Zero);
+
+    // select a, b, !a -> select !a, 1, b
+    // select a, b, zext(!a) -> select !a, 1, b
+    if (match(FalseVal, m_ZExtOrSelf(m_CombineAnd(m_Value(NotCond),
+                                                  m_Not(m_Specific(CondVal))))))
+      return SelectInst::Create(NotCond, One, TrueVal);
+
+    // select a, b, sext(!a) -> select !a, -1, b
+    if (match(FalseVal, m_SExt(m_CombineAnd(m_Value(NotCond),
+                                            m_Not(m_Specific(CondVal))))))
+      return SelectInst::Create(NotCond, AllOnes, TrueVal);
+  }
+
   // Avoid potential infinite loops by checking for non-constant condition.
   // TODO: Can we assert instead by improving canonicalizeSelectToShuffle()?
   //       Scalar select must have simplified?
@@ -2711,20 +2756,6 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
       return SelectInst::Create(NotCond, One, TrueVal);
     }
 
-    // select a, a, b -> select a, true, b
-    if (CondVal == TrueVal)
-      return replaceOperand(SI, 1, One);
-    // select a, b, a -> select a, b, false
-    if (CondVal == FalseVal)
-      return replaceOperand(SI, 2, Zero);
-
-    // select a, !a, b -> select !a, b, false
-    if (match(TrueVal, m_Not(m_Specific(CondVal))))
-      return SelectInst::Create(TrueVal, FalseVal, Zero);
-    // select a, b, !a -> select !a, true, b
-    if (match(FalseVal, m_Not(m_Specific(CondVal))))
-      return SelectInst::Create(FalseVal, One, TrueVal);
-
     Value *A, *B;
 
     // DeMorgan in select form: !a && !b --> !(a || b)

llvm/test/Transforms/InstCombine/select.ll

@@ -3213,8 +3213,7 @@ declare i32 @llvm.cttz.i32(i32, i1 immarg)
 
 define i32 @select_cond_zext_cond(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_zext_cond(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i1 [[COND:%.*]] to i32
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 1, i32 [[B:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[ZEXT]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %zext = zext i1 %cond to i32
@@ -3224,8 +3223,7 @@ define i32 @select_cond_zext_cond(i1 %cond, i32 %b) {
 
 define <2 x i32> @select_cond_zext_cond_vec(<2 x i1> %cond, <2 x i32> %b) {
 ; CHECK-LABEL: @select_cond_zext_cond_vec(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext <2 x i1> [[COND:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> <i32 1, i32 1>, <2 x i32> [[B:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x i32> [[ZEXT]], <2 x i32> [[B:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %zext = zext <2 x i1> %cond to <2 x i32>
@@ -3235,8 +3233,7 @@ define <2 x i32> @select_cond_zext_cond_vec(<2 x i1> %cond, <2 x i32> %b) {
 
 define i32 @select_cond_sext_cond(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_sext_cond(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[COND:%.*]] to i32
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 -1, i32 [[B:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[SEXT]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %sext = sext i1 %cond to i32
@@ -3246,8 +3243,7 @@ define i32 @select_cond_sext_cond(i1 %cond, i32 %b) {
 
 define <2 x i32> @select_cond_sext_cond_vec(<2 x i1> %cond, <2 x i32> %b) {
 ; CHECK-LABEL: @select_cond_sext_cond_vec(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i1> [[COND:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> <i32 -1, i32 -1>, <2 x i32> [[B:%.*]]
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x i32> [[SEXT]], <2 x i32> [[B:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %sext = sext <2 x i1> %cond to <2 x i32>
@@ -3257,8 +3253,7 @@ define <2 x i32> @select_cond_sext_cond_vec(<2 x i1> %cond, <2 x i32> %b) {
 
 define i32 @select_cond_val_zext_cond(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_val_zext_cond(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i1 [[COND:%.*]] to i32
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 0
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[B:%.*]], i32 [[ZEXT]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %zext = zext i1 %cond to i32
@@ -3268,8 +3263,7 @@ define i32 @select_cond_val_zext_cond(i1 %cond, i32 %b) {
 
 define <2 x i32> @select_cond_val_zext_cond_vec(<2 x i1> %cond, <2 x i32> %b) {
 ; CHECK-LABEL: @select_cond_val_zext_cond_vec(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext <2 x i1> [[COND:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> [[B:%.*]], <2 x i32> zeroinitializer
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x i32> [[B:%.*]], <2 x i32> [[ZEXT]]
 ; CHECK-NEXT:    ret <2 x i32> [[SEL]]
 ;
   %zext = zext <2 x i1> %cond to <2 x i32>
@@ -3279,8 +3273,7 @@ define <2 x i32> @select_cond_val_zext_cond_vec(<2 x i1> %cond, <2 x i32> %b) {
 
 define i32 @select_cond_val_sext_cond(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_val_sext_cond(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[COND:%.*]] to i32
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 0
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[B:%.*]], i32 [[SEXT]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %sext = sext i1 %cond to i32
@@ -3290,9 +3283,7 @@ define i32 @select_cond_val_sext_cond(i1 %cond, i32 %b) {
 
 define i32 @select_cond_zext_not_cond_val(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_zext_not_cond_val(
-; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND:%.*]], true
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 0, i32 [[B:%.*]]
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i1 [[NOT_COND]] to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[ZEXT]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %not_cond = xor i1 %cond, true
@@ -3303,9 +3294,7 @@ define i32 @select_cond_zext_not_cond_val(i1 %cond, i32 %b) {
 
 define i32 @select_cond_sext_not_cond_val(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_sext_not_cond_val(
-; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND:%.*]], true
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 0, i32 [[B:%.*]]
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[NOT_COND]] to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[SEXT]], i32 [[B:%.*]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %not_cond = xor i1 %cond, true
@@ -3317,9 +3306,7 @@ define i32 @select_cond_sext_not_cond_val(i1 %cond, i32 %b) {
 
 define i32 @select_cond_val_zext_not_cond(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_val_zext_not_cond(
-; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND:%.*]], true
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 1
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i1 [[NOT_COND]] to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[B:%.*]], i32 [[ZEXT]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %not_cond = xor i1 %cond, true
@@ -3330,9 +3317,7 @@ define i32 @select_cond_val_zext_not_cond(i1 %cond, i32 %b) {
 
 define i32 @select_cond_val_sext_not_cond(i1 %cond, i32 %b) {
 ; CHECK-LABEL: @select_cond_val_sext_not_cond(
-; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND:%.*]], true
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 -1
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[NOT_COND]] to i32
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], i32 [[B:%.*]], i32 [[SEXT]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %not_cond = xor i1 %cond, true
@@ -3343,11 +3328,7 @@ define i32 @select_cond_val_sext_not_cond(i1 %cond, i32 %b) {
 
 define i32 @select_cond_not_cond_cond1(i1 %cond) {
 ; CHECK-LABEL: @select_cond_not_cond_cond1(
-; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[COND:%.*]] to i32
+; CHECK-NEXT:    ret i32 0
-; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND]], true
-; CHECK-NEXT:    [[S:%.*]] = sext i1 [[NOT_COND]] to i32
-; CHECK-NEXT:    [[V:%.*]] = select i1 [[COND]], i32 [[S]], i32 [[Z]]
-; CHECK-NEXT:    ret i32 [[V]]
 ;
   %z = zext i1 %cond to i32
   %not_cond = xor i1 %cond, true
@@ -3358,11 +3339,7 @@ define i32 @select_cond_not_cond_cond1(i1 %cond) {
 
 define i32 @select_cond_not_cond_cond2(i1 %cond) {
 ; CHECK-LABEL: @select_cond_not_cond_cond2(
-; CHECK-NEXT:    [[Z:%.*]] = sext i1 [[COND:%.*]] to i32
+; CHECK-NEXT:    ret i32 0
-; CHECK-NEXT:    [[NOT_COND:%.*]] = xor i1 [[COND]], true
-; CHECK-NEXT:    [[S:%.*]] = zext i1 [[NOT_COND]] to i32
-; CHECK-NEXT:    [[V:%.*]] = select i1 [[COND]], i32 [[S]], i32 [[Z]]
-; CHECK-NEXT:    ret i32 [[V]]
 ;
   %z = sext i1 %cond to i32
   %not_cond = xor i1 %cond, true