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

[zero9178]

Summary

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

[RKSimon]

llvm/test/Transforms/InstSimplify/select.ll

1121 ↗

(On Diff #461090)

please can you add vector test coverage? and possibly a negative test for the case where the true/false args are commuted?

[zero9178]

Address review comments: Add test using vector types and a negative test with switched true false operands

[spatel]

This match seems over-specified. If either the true or false arm of the select are cast from the select condition, then we can replace that arm in -instcombine, and it should collapse to zero when both arms match:
https://alive2.llvm.org/ce/z/35kadA

[zero9178]

Address review comments: Change approach to handling it in InstCombine. I noticed that these "implication" patterns for the condition being used as either the true or false values already existed but have simply only been handled for i1 selection type while not taking zext and sext into account.
I now made these also handle zext and sext as well, making it capable of handling any integer type.

[spatel]

In D134142#3802487, @zero9178 wrote:

Address review comments: Change approach to handling it in InstCombine. I noticed that these "implication" patterns for the condition being used as either the true or false values already existed but have simply only been handled for i1 selection type while not taking zext and sext into account.
I now made these also handle zext and sext as well, making it capable of handling any integer type.

Nice improvement! Do you want to add a couple of tests to model your motivating example (where both arms of the select are extended)? Please pre-commit the tests with baseline output, then update the patch here, so it shows test diffs.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2671–2672

Nit: I'd make all of these ConstantInt::get***() for consistency since we know all of them are integer types.

2683–2685

Change the "false" to "0" to be more general in these comments (and similar for "true" -> "1" below).

[nikic]

This kinda looks like foldSelectValueEquivalence() with a trivial comparison (i.e. interpret c as c == 1).

[zero9178]

Address review comments:

• Consistently use ConstantInt
• Use 0 and 1 instead of true and false in comments, since it is now applicable to all integer types, not just i1
• Add test coverage for original motivating example. I only added the zext/sext and sext/zext cases, as the other cases are comparatively uniniteresting and were previously already handled by sinking of the casts below the select
• Precommit tests and rebase this patch on top of it

[spatel]

In D134142#3802599, @nikic wrote:

This kinda looks like foldSelectValueEquivalence() with a trivial comparison (i.e. interpret c as c == 1).

Agree - but we'd have to account for the missing icmp that code is expecting? Not sure that's worth the effort and reduced readability, so this patch LGTM.

[nikic]

In D134142#3802765, @spatel wrote:

In D134142#3802599, @nikic wrote:

This kinda looks like foldSelectValueEquivalence() with a trivial comparison (i.e. interpret c as c == 1).

Agree - but we'd have to account for the missing icmp that code is expecting? Not sure that's worth the effort and reduced readability, so this patch LGTM.

It looks like we don't even need foldSelectValueEquivalence here -- we already have simplifyWithOpReplaced() folds for these. They are currently behind a guard for i1 selects, but if we move them outside we get most of the mileage: https://gist.github.com/nikic/86420628e72dd137aa56d2317f3825b0

The main caveat is that this does not handle vectors. Not that I actually care about vectors, but for the sake of propriety we should probably update simplifyWithOpReplaced() to allow vectors and only reject the problematic cross-lane operations inside.

[spatel]

In D134142#3802960, @nikic wrote:

In D134142#3802765, @spatel wrote:

In D134142#3802599, @nikic wrote:

This kinda looks like foldSelectValueEquivalence() with a trivial comparison (i.e. interpret c as c == 1).

Agree - but we'd have to account for the missing icmp that code is expecting? Not sure that's worth the effort and reduced readability, so this patch LGTM.

It looks like we don't even need foldSelectValueEquivalence here -- we already have simplifyWithOpReplaced() folds for these. They are currently behind a guard for i1 selects, but if we move them outside we get most of the mileage: https://gist.github.com/nikic/86420628e72dd137aa56d2317f3825b0

The main caveat is that this does not handle vectors. Not that I actually care about vectors, but for the sake of propriety we should probably update simplifyWithOpReplaced() to allow vectors and only reject the problematic cross-lane operations inside.

That also misses the basic patterns with a 'not' before the extend. I like the smaller patch, but given that this one covers more cases, I'd still go with this patch as-is.

[nikic]

Okay, let's go with this for now. I can come back later and refactor this to use generic folds when I have time.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (Value *V = simplifySelectInst(CondVal, TrueVal, FalseVal,
     return replaceInstUsesWith(SI, V);
 
   if (Instruction *I = canonicalizeSelectToShuffle(SI))
     return I;
 
   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 = Constant::getNullValue(SelType);
+    auto *AllOnes = Constant::getAllOnesValue(SelType);
+
+    // 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, false
+    // select a, b, zext(a) -> select a, b, false
+    // select a, b, sext(a) -> select a, b, false
+    if (match(FalseVal, m_ZExtOrSExtOrSelf(m_Specific(CondVal))))
+      return replaceOperand(SI, 2, Zero);
+
+    Value *NotCond;
+
+    // select a, !a, b -> select !a, b, false
+    // select a, sext(!a), b -> select !a, b, false
+    // select a, zext(!a), b -> select !a, b, false
+    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, true, b
+    // select a, b, zext(!a) -> select !a, true, 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?
   if (SelType->isIntOrIntVectorTy(1) && !isa<Constant>(CondVal) &&
       TrueVal->getType() == CondVal->getType()) {
     // Folding select to and/or i1 isn't poison safe in general. impliesPoison
     // checks whether folding it does not convert a well-defined value into
     // poison.
if (match(TrueVal, m_Specific(Zero))) {
       return SelectInst::Create(NotCond, FalseVal, Zero);
     }
     // select a, b, true -> select !a, true, b
     if (match(FalseVal, m_Specific(One))) {
       Value *NotCond = Builder.CreateNot(CondVal, "not." + CondVal->getName());
       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)
     // select !a, !b, false --> not (select a, true, b)
     if (match(&SI, m_LogicalAnd(m_Not(m_Value(A)), m_Not(m_Value(B)))) &&
         (CondVal->hasOneUse() || TrueVal->hasOneUse()) &&
         !match(A, m_ConstantExpr()) && !match(B, m_ConstantExpr()))
       return BinaryOperator::CreateNot(Builder.CreateSelect(A, One, B));

llvm/test/Transforms/InstCombine/select.ll

;
   %op = udiv i64 %sel, 3
   ret i64 %op
 }
 
 declare void @use(i1)
 declare void @use_i8(i8)
 declare void @use_i32(i32)
 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:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 1, i32 [[B:%.*]]
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %zext = zext i1 %cond to i32
+  %sel = select i1 %cond, i32 %zext, i32 %b
+  ret i32 %sel
+}
+
+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:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> <i32 1, i32 1>, <2 x i32> [[B:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[SEL]]
+;
+  %zext = zext <2 x i1> %cond to <2 x i32>
+  %sel = select <2 x i1> %cond, <2 x i32> %zext, <2 x i32> %b
+  ret <2 x i32> %sel
+}
+
+define i32 @select_cond_sext_cond(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_sext_cond(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 -1, i32 [[B:%.*]]
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %sext = sext i1 %cond to i32
+  %sel = select i1 %cond, i32 %sext, i32 %b
+  ret i32 %sel
+}
+
+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:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> <i32 -1, i32 -1>, <2 x i32> [[B:%.*]]
+; CHECK-NEXT:    ret <2 x i32> [[SEL]]
+;
+  %sext = sext <2 x i1> %cond to <2 x i32>
+  %sel = select <2 x i1> %cond, <2 x i32> %sext, <2 x i32> %b
+  ret <2 x i32> %sel
+}
+
+define i32 @select_cond_val_zext_cond(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_val_zext_cond(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 0
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %zext = zext i1 %cond to i32
+  %sel = select i1 %cond, i32 %b, i32 %zext
+  ret i32 %sel
+}
+
+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:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> [[B:%.*]], <2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <2 x i32> [[SEL]]
+;
+  %zext = zext <2 x i1> %cond to <2 x i32>
+  %sel = select <2 x i1> %cond, <2 x i32> %b, <2 x i32> %zext
+  ret <2 x i32> %sel
+}
+
+define i32 @select_cond_val_sext_cond(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_val_sext_cond(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 0
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %sext = sext i1 %cond to i32
+  %sel = select i1 %cond, i32 %b, i32 %sext
+  ret i32 %sel
+}
+
+define i32 @select_cond_zext_not_cond_val(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_zext_not_cond_val(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 0, i32 [[B:%.*]]
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %not_cond = xor i1 %cond, true
+  %zext = zext i1 %not_cond to i32
+  %sel = select i1 %cond, i32 %zext, i32 %b
+  ret i32 %sel
+}
+
+define i32 @select_cond_sext_not_cond_val(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_sext_not_cond_val(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 0, i32 [[B:%.*]]
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %not_cond = xor i1 %cond, true
+  %sext = sext i1 %not_cond to i32
+  %sel = select i1 %cond, i32 %sext, i32 %b
+  ret i32 %sel
+}
+
+
+define i32 @select_cond_val_zext_not_cond(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_val_zext_not_cond(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 1
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %not_cond = xor i1 %cond, true
+  %zext = zext i1 %not_cond to i32
+  %sel = select i1 %cond, i32 %b, i32 %zext
+  ret i32 %sel
+}
+
+define i32 @select_cond_val_sext_not_cond(i1 %cond, i32 %b) {
+; CHECK-LABEL: @select_cond_val_sext_not_cond(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 [[B:%.*]], i32 -1
+; CHECK-NEXT:    ret i32 [[SEL]]
+;
+  %not_cond = xor i1 %cond, true
+  %sext = sext i1 %not_cond to i32
+  %sel = select i1 %cond, i32 %b, i32 %sext
+  ret i32 %sel
+}
+
+