[InstCombine]: foldSelectICmpAndAnd(): and is commutative

[lebedev.ri]

Summary

The fold added in D45108 did not account for the fact that
the and instruction is commutative, and if the mask is a variable,
the mask variable and the fold variable may be swapped.

I have noticed this by accident when looking into PR6773

This extends/generalizes that fold, so it is handled too.

[spatel]

I think the patch is correct, but it's hard to keep all of the values straight (because there are a lot!). What do you think about rearranging it a bit like this and using m_c_And():

/// We want to turn:
///   (select (icmp eq (and X, Y), 0), (and (lshr X, Z), 1), 1)
/// into:
///   zext (icmp ne i32 (and X, (or Y, (shl 1, Z))), 0)
/// Note:
///   Z may be 0 if lshr is missing.
/// Worst-case scenario is that we will replace 5 instructions with 5 different
/// instructions, but we got rid of select.
static Instruction *foldSelectICmpAndAnd(Type *SelType, const ICmpInst *Cmp,
                                         Value *TVal, Value *FVal,
                                         InstCombiner::BuilderTy &Builder) {
  if (!Cmp->hasOneUse() || !Cmp->getOperand(0)->hasOneUse() ||
      Cmp->getPredicate() != ICmpInst::ICMP_EQ ||
      !match(Cmp->getOperand(1), m_Zero()) || !match(FVal, m_One()))
    return nullptr;

  // The TrueVal has general form of:
  //   and %B, 1
  // where %B may be optionally shifted: lshr %X, %Z.
  Value *B;
  if (!match(TVal, m_OneUse(m_And(m_Value(B), m_One()))))
    return nullptr;

  Value *X, *Y, *Z;
  bool HasShift = match(B, m_OneUse(m_LShr(m_Value(X), m_Value(Z))));
  if (!HasShift)
    X = B;

  if (!match(Cmp->getOperand(0), m_c_And(m_Specific(X), m_Value(Y))))
    return nullptr;

  // ((X & Y) == 0) ? ((X >> Z) & 1) : 1 --> (X & (Y | (1 << Z))) != 0
  // ((X & Y) == 0) ? (X & 1) : 1 --> (X & (Y | 1)) != 0
  Constant *One = ConstantInt::get(SelType, 1);
  Value *MaskB = HasShift ? Builder.CreateShl(One, Z) : One;
  Value *FullMask = Builder.CreateOr(Y, MaskB);
  Value *MaskedX = Builder.CreateAnd(X, FullMask);
  Value *ICmpNeZero = Builder.CreateIsNotNull(MaskedX);
  return new ZExtInst(ICmpNeZero, SelType);
}

[spatel]

Side note: this patch / commit should not be labeled 'NFC'.

[lebedev.ri]

In D45539#1066245, @spatel wrote:

Side note: this patch / commit should not be labeled 'NFC'.

Whoops, too greedy copypasting from the D45538 subject, sorry!

[lebedev.ri]

Make foldSelectICmpAndAnd() much smaller, as suggested by @spatel.

[spatel]

LGTM.

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

 }
 
 /// We want to turn:
 ///   (select (icmp eq (and X, Y), 0), (and (lshr X, Z), 1), 1)
 /// into:
 ///   zext (icmp ne i32 (and X, (or Y, (shl 1, Z))), 0)
 /// Note:
 ///   Z may be 0 if lshr is missing.
-/// Worst case scenario is that we will replace 5 instructions with 5 different
+/// Worst-case scenario is that we will replace 5 instructions with 5 different
 /// instructions, but we got rid of select.
-static Instruction *foldSelectICmpAndAnd(Type *SelType, const ICmpInst *IC,
-                                         Value *TrueVal, Value *FalseVal,
+static Instruction *foldSelectICmpAndAnd(Type *SelType, const ICmpInst *Cmp,
+                                         Value *TVal, Value *FVal,
                                          InstCombiner::BuilderTy &Builder) {
-  if (!(IC->hasOneUse() && IC->getOperand(0)->hasOneUse()))
+  if (!(Cmp->hasOneUse() && Cmp->getOperand(0)->hasOneUse() &&
+        Cmp->getPredicate() == ICmpInst::ICMP_EQ &&
+        match(Cmp->getOperand(1), m_Zero()) && match(FVal, m_One())))
     return nullptr;
 
-  Value *X, *Y;
-  ICmpInst::Predicate EqPred;
-  if (!(match(IC, m_ICmp(EqPred, m_And(m_Value(X), m_Value(Y)), m_Zero())) &&
-        ICmpInst::Predicate::ICMP_EQ == EqPred && match(FalseVal, m_One())))
-    return nullptr;
-
-  // The TrueVal has general form of:
-  //   and %B, 1
+  // The TrueVal has general form of:  and %B, 1
   Value *B;
-  if (!match(TrueVal, m_OneUse(m_And(m_Value(B), m_One()))))
+  if (!match(TVal, m_OneUse(m_And(m_Value(B), m_One()))))
     return nullptr;
 
-  // Where %B can be one of:
-  //        %X
-  // or
-  //   lshr %X, %Z
-  // Where %Z may or may not be a constant.
-  Value *MaskB, *Z;
-  if (match(B, m_Specific(X))) {
-    MaskB = ConstantInt::get(SelType, 1);
-  } else if (match(B, m_OneUse(m_LShr(m_Specific(X), m_Value(Z))))) {
-    MaskB = Builder.CreateShl(ConstantInt::get(SelType, 1), Z);
-  } else
+  // Where %B may be optionally shifted:  lshr %X, %Z.
+  Value *X, *Z;
+  const bool HasShift = match(B, m_OneUse(m_LShr(m_Value(X), m_Value(Z))));
+  if (!HasShift)
+    X = B;
+
+  Value *Y;
+  if (!match(Cmp->getOperand(0), m_c_And(m_Specific(X), m_Value(Y))))
     return nullptr;
 
+  // ((X & Y) == 0) ? ((X >> Z) & 1) : 1 --> (X & (Y | (1 << Z))) != 0
+  // ((X & Y) == 0) ? (X & 1) : 1 --> (X & (Y | 1)) != 0
+  Constant *One = ConstantInt::get(SelType, 1);
+  Value *MaskB = HasShift ? Builder.CreateShl(One, Z) : One;
   Value *FullMask = Builder.CreateOr(Y, MaskB);
   Value *MaskedX = Builder.CreateAnd(X, FullMask);
   Value *ICmpNeZero = Builder.CreateIsNotNull(MaskedX);
   return new ZExtInst(ICmpNeZero, SelType);
 }
 
 /// We want to turn:
 ///   (select (icmp eq (and X, C1), 0), Y, (or Y, C2))

llvm/trunk/test/Transforms/InstCombine/select-of-bittest.ll

;
   %tmp4 = and i32 %tmp3, 1
   %tmp5 = select i1 %tmp2, i32 %tmp4, i32 1
   ret i32 %tmp5
 }
 
 ; Should be exactly as the previous one
 define i32 @f_var0_commutative_and(i32 %arg, i32 %arg1) {
 ; CHECK-LABEL: @f_var0_commutative_and(
-; CHECK-NEXT:    [[TMP:%.*]] = and i32 [[ARG1:%.*]], [[ARG:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP]], 0
-; CHECK-NEXT:    [[TMP3:%.*]] = lshr i32 [[ARG]], 1
-; CHECK-NEXT:    [[TMP4:%.*]] = and i32 [[TMP3]], 1
-; CHECK-NEXT:    [[TMP5:%.*]] = select i1 [[TMP2]], i32 [[TMP4]], i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[ARG1:%.*]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], [[ARG:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne i32 [[TMP2]], 0
+; CHECK-NEXT:    [[TMP5:%.*]] = zext i1 [[TMP3]] to i32
 ; CHECK-NEXT:    ret i32 [[TMP5]]
 ;
   %tmp = and i32 %arg1, %arg ; in different order
   %tmp2 = icmp eq i32 %tmp, 0
   %tmp3 = lshr i32 %arg, 1
   %tmp4 = and i32 %tmp3, 1
   %tmp5 = select i1 %tmp2, i32 %tmp4, i32 1
   ret i32 %tmp5
;
   %tmp3 = and i32 %arg, 1
   %tmp4 = select i1 %tmp2, i32 %tmp3, i32 1
   ret i32 %tmp4
 }
 
 ; Should be exactly as the previous one
 define i32 @f_var1_commutative_and(i32 %arg, i32 %arg1) {
 ; CHECK-LABEL: @f_var1_commutative_and(
-; CHECK-NEXT:    [[TMP:%.*]] = and i32 [[ARG1:%.*]], [[ARG:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP]], 0
-; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[ARG]], 1
-; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i32 [[TMP3]], i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[ARG1:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], [[ARG:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne i32 [[TMP2]], 0
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i1 [[TMP3]] to i32
 ; CHECK-NEXT:    ret i32 [[TMP4]]
 ;
   %tmp = and i32 %arg1, %arg ; in different order
   %tmp2 = icmp eq i32 %tmp, 0
   %tmp3 = and i32 %arg, 1
   %tmp4 = select i1 %tmp2, i32 %tmp3, i32 1
   ret i32 %tmp4
 }
;
   %tmp5 = and i32 %tmp4, 1
   %tmp6 = select i1 %tmp3, i32 %tmp5, i32 1
   ret i32 %tmp6
 }
 
 ; Should be exactly as the previous one
 define i32 @f_var3_commutative_and(i32 %arg, i32 %arg1, i32 %arg2) {
 ; CHECK-LABEL: @f_var3_commutative_and(
-; CHECK-NEXT:    [[TMP:%.*]] = and i32 [[ARG1:%.*]], [[ARG:%.*]]
-; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP]], 0
-; CHECK-NEXT:    [[TMP4:%.*]] = lshr i32 [[ARG]], [[ARG2:%.*]]
-; CHECK-NEXT:    [[TMP5:%.*]] = and i32 [[TMP4]], 1
-; CHECK-NEXT:    [[TMP6:%.*]] = select i1 [[TMP3]], i32 [[TMP5]], i32 1
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 1, [[ARG2:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = or i32 [[TMP1]], [[ARG1:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], [[ARG:%.*]]
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[TMP3]], 0
+; CHECK-NEXT:    [[TMP6:%.*]] = zext i1 [[TMP4]] to i32
 ; CHECK-NEXT:    ret i32 [[TMP6]]
 ;
   %tmp = and i32 %arg1, %arg ; in different order
   %tmp3 = icmp eq i32 %tmp, 0
   %tmp4 = lshr i32 %arg, %arg2
   %tmp5 = and i32 %tmp4, 1
   %tmp6 = select i1 %tmp3, i32 %tmp5, i32 1
   ret i32 %tmp6