[InstCombine] extend matchSelectFromAndOr() to work with i1 scalar types

[spatel]

Summary

If the incoming types are i1, then we don't have to pattern match any sext ops.

[eli.friedman]

Would it make sense to add "m_SExtOrBool" or something like that, which matches a value which matches either a sext or a value of boolean type? That seems like a cleaner way to generalize the transform.

If the boolean comes from an icmp, it seems likely that you'll end up with ((X == Y) & C) | ((X != Y) & D). I guess you can address that in a followup, if necessary.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1654 ↗

(On Diff #61921)

Please don't get rid of this comment... you don't have any others which actually describe the transform in full.

[spatel]

In D21740#470799, @eli.friedman wrote:

Would it make sense to add "m_SExtOrBool" or something like that, which matches a value which matches either a sext or a value of boolean type? That seems like a cleaner way to generalize the transform.

Not sure if that matcher would be generally useful. For the non-splat vectors scenarios that I'm seeing, I may need to add something to ValueTracking like 'isAllOnesOrAllZerosElts()' [return true if a scalar value is -1 or 0 or if all elements of a vector are -1 or 0 (but not necessarily that all elements are the same). So that would be an even more general way to catch these. I'd prefer to keep this patch simpler and then see what's needed for the 'TODO' comment.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1654 ↗

(On Diff #61921)

Oops - that was an unintended delete. Let me restore it.

[spatel]

Patch updated:
Restore mistakenly deleted comment that describes the transform.

[eli.friedman]

LGTM.

llvm/trunk/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Instruction *InstCombiner::MatchBSwap(BinaryOperator &I) {
   Instruction *LastInst = Insts.pop_back_val();
   LastInst->removeFromParent();
 
   for (auto *Inst : Insts)
     Worklist.Add(Inst);
   return LastInst;
 }
 
-/// We have an expression of the form (A & C) | (B & D). If A is (Cond?-1:0)
-/// and B is ~(Cond?-1,0), then simplify this expression to "Cond ? C : D".
+/// We have an expression of the form (A & C) | (B & D). If A is a scalar or
+/// vector composed of all-zeros or all-ones values and is the bitwise 'not' of
+/// B, it can be used as the condition operand of a select instruction.
+static Value *getSelectCondition(Value *A, Value *B) {
+  // If these are scalars or vectors of i1, A can be used directly.
+  Type *Ty = A->getType();
+  if (match(A, m_Not(m_Specific(B))) && Ty->getScalarType()->isIntegerTy(1))
+    return A;
+
+  // If A and B are sign-extended, look through the sexts to find the booleans.
+  Value *Cond;
+  if (match(A, m_SExt(m_Value(Cond))) &&
+      Cond->getType()->getScalarType()->isIntegerTy(1) &&
+      match(B, m_CombineOr(m_Not(m_SExt(m_Specific(Cond))),
+                           m_SExt(m_Not(m_Specific(Cond))))))
+    return Cond;
+
+  // TODO: Try more matches that only apply to non-splat constant vectors.
+
+  return nullptr;
+}
+
+/// We have an expression of the form (A & C) | (B & D). Try to simplify this
+/// to "A' ? C : D", where A' is a boolean or vector of booleans.
 static Value *matchSelectFromAndOr(Value *A, Value *C, Value *B, Value *D,
-                                         InstCombiner::BuilderTy &Builder) {
   // The potential condition of the select may be bitcasted. In that case, look
   // through its bitcast and the corresponding bitcast of the 'not' condition.
   Type *OrigType = A->getType();
   Value *SrcA, *SrcB;
   if (match(A, m_BitCast(m_Value(SrcA))) &&
       match(B, m_BitCast(m_Value(SrcB)))) {
     A = SrcA;
     B = SrcB;
   }
 
-  // The condition must be a value of -1/0, and B must be the 'not' of that
-  // condition.
-  Value *Cond;
-  if (match(A, m_SExt(m_Value(Cond))) &&
-      Cond->getType()->getScalarType()->isIntegerTy(1) &&
-      match(B, m_CombineOr(m_Not(m_SExt(m_Specific(Cond))),
-                           m_SExt(m_Not(m_Specific(Cond)))))) {
+  if (Value *Cond = getSelectCondition(A, B)) {
     // ((bc Cond) & C) | ((bc ~Cond) & D) --> bc (select Cond, (bc C), (bc D))
     // The bitcasts will either all exist or all not exist. The builder will
     // not create unnecessary casts if the types already match.
     Value *BitcastC = Builder.CreateBitCast(C, A->getType());
     Value *BitcastD = Builder.CreateBitCast(D, A->getType());
     Value *Select = Builder.CreateSelect(Cond, BitcastC, BitcastD);
     return Builder.CreateBitCast(Select, OrigType);
   }

llvm/trunk/test/Transforms/InstCombine/logical-select.ll

;
   %and1 = and <2 x i64> %sia, %bc1
   %neg = xor <4 x i32> %sext, <i32 -1, i32 -1, i32 -1, i32 -1>
   %bc2 = bitcast <4 x i32> %neg to <2 x i64>
   %and2 = and <2 x i64> %sib, %bc2
   %or = or <2 x i64> %and2, %and1
   ret <2 x i64> %or
 }
 
-; FIXME: Missed conversions to select below here.
-
 define i1 @bools(i1 %a, i1 %b, i1 %c) {
 ; CHECK-LABEL: @bools(
-; CHECK-NEXT:    [[NOT:%.*]] = xor i1 %c, true
-; CHECK-NEXT:    [[AND1:%.*]] = and i1 [[NOT]], %a
-; CHECK-NEXT:    [[AND2:%.*]] = and i1 %c, %b
-; CHECK-NEXT:    [[OR:%.*]] = or i1 [[AND1]], [[AND2]]
-; CHECK-NEXT:    ret i1 [[OR]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 %c, i1 %b, i1 %a
+; CHECK-NEXT:    ret i1 [[TMP1]]
 ;
   %not = xor i1 %c, -1
   %and1 = and i1 %not, %a
   %and2 = and i1 %c, %b
   %or = or i1 %and1, %and2
   ret i1 %or
 }
 
 define <4 x i1> @vec_of_bools(<4 x i1> %a, <4 x i1> %b, <4 x i1> %c) {
 ; CHECK-LABEL: @vec_of_bools(
-; CHECK-NEXT:    [[NOT:%.*]] = xor <4 x i1> %c, <i1 true, i1 true, i1 true, i1 true>
-; CHECK-NEXT:    [[AND1:%.*]] = and <4 x i1> [[NOT]], %a
-; CHECK-NEXT:    [[AND2:%.*]] = and <4 x i1> %b, %c
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i1> [[AND2]], [[AND1]]
-; CHECK-NEXT:    ret <4 x i1> [[OR]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select <4 x i1> %c, <4 x i1> %b, <4 x i1> %a
+; CHECK-NEXT:    ret <4 x i1> [[TMP1]]
 ;
   %not = xor <4 x i1> %c, <i1 true, i1 true, i1 true, i1 true>
   %and1 = and <4 x i1> %not, %a
   %and2 = and <4 x i1> %b, %c
   %or = or <4 x i1> %and2, %and1
   ret <4 x i1> %or
 }
 
 define i4 @vec_of_casted_bools(i4 %a, i4 %b, <4 x i1> %c) {
 ; CHECK-LABEL: @vec_of_casted_bools(
-; CHECK-NEXT:    [[NOT:%.*]] = xor <4 x i1> %c, <i1 true, i1 true, i1 true, i1 true>
-; CHECK-NEXT:    [[BC1:%.*]] = bitcast <4 x i1> [[NOT]] to i4
-; CHECK-NEXT:    [[BC2:%.*]] = bitcast <4 x i1> %c to i4
-; CHECK-NEXT:    [[AND1:%.*]] = and i4 [[BC1]], %a
-; CHECK-NEXT:    [[AND2:%.*]] = and i4 [[BC2]], %b
-; CHECK-NEXT:    [[OR:%.*]] = or i4 [[AND1]], [[AND2]]
-; CHECK-NEXT:    ret i4 [[OR]]
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i4 %a to <4 x i1>
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i4 %b to <4 x i1>
+; CHECK-NEXT:    [[TMP3:%.*]] = select <4 x i1> %c, <4 x i1> [[TMP2]], <4 x i1> [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast <4 x i1> [[TMP3]] to i4
+; CHECK-NEXT:    ret i4 [[TMP4]]
 ;
   %not = xor <4 x i1> %c, <i1 true, i1 true, i1 true, i1 true>
   %bc1 = bitcast <4 x i1> %not to i4
   %bc2 = bitcast <4 x i1> %c to i4
   %and1 = and i4 %a, %bc1
   %and2 = and i4 %bc2, %b
   %or = or i4 %and1, %and2
   ret i4 %or
 }
 
+; FIXME: Missed conversions to select below here.
 ; Inverted 'and' constants mean this is a select.
 
 define <4 x i32> @vec_sel_consts(<4 x i32> %a, <4 x i32> %b) {
 ; CHECK-LABEL: @vec_sel_consts(
 ; CHECK-NEXT:    [[AND1:%.*]] = and <4 x i32> %a, <i32 -1, i32 0, i32 0, i32 -1>
 ; CHECK-NEXT:    [[AND2:%.*]] = and <4 x i32> %b, <i32 0, i32 -1, i32 -1, i32 0>
 ; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[AND1]], [[AND2]]
 ; CHECK-NEXT:    ret <4 x i32> [[OR]]