[InstCombine] Support logical and in masked icmp fold

Most of the folds implemented in this function work fine with
logical operations. We only need to be careful for the cases that
work on non-constant masks, where the RHS operand shouldn't be
poison.

This is a conservative implementation that bails out of illegal
transforms, but we could also change these to insert freeze instead.

Author: nikic

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -365,6 +365,7 @@ getMaskedTypeForICmpPair(Value *&A, Value *&B, Value *&C,
 /// (icmp(A & X) ==/!= Y), where the left-hand side is of type Mask_NotAllZeros
 /// and the right hand side is of type BMask_Mixed. For example,
 /// (icmp (A & 12) != 0) & (icmp (A & 15) == 8) -> (icmp (A & 15) == 8).
+/// Also used for logical and/or, must be poison safe.
 static Value *foldLogOpOfMaskedICmps_NotAllZeros_BMask_Mixed(
     ICmpInst *LHS, ICmpInst *RHS, bool IsAnd, Value *A, Value *B, Value *C,
     Value *D, Value *E, ICmpInst::Predicate PredL, ICmpInst::Predicate PredR,
@@ -486,6 +487,7 @@ static Value *foldLogOpOfMaskedICmps_NotAllZeros_BMask_Mixed(
 /// Try to fold (icmp(A & B) ==/!= 0) &/| (icmp(A & D) ==/!= E) into a single
 /// (icmp(A & X) ==/!= Y), where the left-hand side and the right hand side
 /// aren't of the common mask pattern type.
+/// Also used for logical and/or, must be poison safe.
 static Value *foldLogOpOfMaskedICmpsAsymmetric(
     ICmpInst *LHS, ICmpInst *RHS, bool IsAnd, Value *A, Value *B, Value *C,
     Value *D, Value *E, ICmpInst::Predicate PredL, ICmpInst::Predicate PredR,
@@ -520,6 +522,7 @@ static Value *foldLogOpOfMaskedICmpsAsymmetric(
 /// Try to fold (icmp(A & B) ==/!= C) &/| (icmp(A & D) ==/!= E)
 /// into a single (icmp(A & X) ==/!= Y).
 static Value *foldLogOpOfMaskedICmps(ICmpInst *LHS, ICmpInst *RHS, bool IsAnd,
+                                     bool IsLogical,
                                      InstCombiner::BuilderTy &Builder) {
   Value *A = nullptr, *B = nullptr, *C = nullptr, *D = nullptr, *E = nullptr;
   ICmpInst::Predicate PredL = LHS->getPredicate(), PredR = RHS->getPredicate();
@@ -564,6 +567,8 @@ static Value *foldLogOpOfMaskedICmps(ICmpInst *LHS, ICmpInst *RHS, bool IsAnd,
   if (Mask & Mask_AllZeros) {
     // (icmp eq (A & B), 0) & (icmp eq (A & D), 0)
     // -> (icmp eq (A & (B|D)), 0)
+    if (IsLogical && !isGuaranteedNotToBeUndefOrPoison(D))
+      return nullptr; // TODO: Use freeze?
     Value *NewOr = Builder.CreateOr(B, D);
     Value *NewAnd = Builder.CreateAnd(A, NewOr);
     // We can't use C as zero because we might actually handle
@@ -575,13 +580,17 @@ static Value *foldLogOpOfMaskedICmps(ICmpInst *LHS, ICmpInst *RHS, bool IsAnd,
   if (Mask & BMask_AllOnes) {
     // (icmp eq (A & B), B) & (icmp eq (A & D), D)
     // -> (icmp eq (A & (B|D)), (B|D))
+    if (IsLogical && !isGuaranteedNotToBeUndefOrPoison(D))
+      return nullptr; // TODO: Use freeze?
     Value *NewOr = Builder.CreateOr(B, D);
     Value *NewAnd = Builder.CreateAnd(A, NewOr);
     return Builder.CreateICmp(NewCC, NewAnd, NewOr);
   }
   if (Mask & AMask_AllOnes) {
     // (icmp eq (A & B), A) & (icmp eq (A & D), A)
     // -> (icmp eq (A & (B&D)), A)
+    if (IsLogical && !isGuaranteedNotToBeUndefOrPoison(D))
+      return nullptr; // TODO: Use freeze?
     Value *NewAnd1 = Builder.CreateAnd(B, D);
     Value *NewAnd2 = Builder.CreateAnd(A, NewAnd1);
     return Builder.CreateICmp(NewCC, NewAnd2, A);
@@ -2442,15 +2451,11 @@ Value *InstCombinerImpl::foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
     }
   }
 
-  // TODO: Some (but not all) of the patterns handled by this function are
-  // safe with logical and/or.
-  if (!IsLogical) {
-    // handle (roughly):
-    // (icmp ne (A & B), C) | (icmp ne (A & D), E)
-    // (icmp eq (A & B), C) & (icmp eq (A & D), E)
-    if (Value *V = foldLogOpOfMaskedICmps(LHS, RHS, IsAnd, Builder))
-      return V;
-  }
+  // handle (roughly):
+  // (icmp ne (A & B), C) | (icmp ne (A & D), E)
+  // (icmp eq (A & B), C) & (icmp eq (A & D), E)
+  if (Value *V = foldLogOpOfMaskedICmps(LHS, RHS, IsAnd, IsLogical, Builder))
+    return V;
 
   // TODO: One of these directions is fine with logical and/or, the other could
   // be supported by inserting freeze.

llvm/test/Transforms/InstCombine/and-or-icmps.ll

@@ -1981,13 +1981,10 @@ define i1 @logical_or_logical_or_icmps_comm3(i8 %x, i8 %y, i8 %z) {
 
 define i1 @bitwise_and_logical_and_masked_icmp_asymmetric(i1 %c, i32 %x) {
 ; CHECK-LABEL: @bitwise_and_logical_and_masked_icmp_asymmetric(
-; CHECK-NEXT:    [[X_M1:%.*]] = and i32 [[X:%.*]], 255
-; CHECK-NEXT:    [[C1:%.*]] = icmp ne i32 [[X_M1]], 0
-; CHECK-NEXT:    [[AND1:%.*]] = select i1 [[C1]], i1 [[C:%.*]], i1 false
-; CHECK-NEXT:    [[X_M2:%.*]] = and i32 [[X]], 11
+; CHECK-NEXT:    [[X_M2:%.*]] = and i32 [[X:%.*]], 11
 ; CHECK-NEXT:    [[C2:%.*]] = icmp eq i32 [[X_M2]], 11
-; CHECK-NEXT:    [[AND2:%.*]] = and i1 [[AND1]], [[C2]]
-; CHECK-NEXT:    ret i1 [[AND2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[C2]], i1 [[C:%.*]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP1]]
 ;
   %x.m1 = and i32 %x, 255
   %c1 = icmp ne i32 %x.m1, 0
@@ -2000,13 +1997,10 @@ define i1 @bitwise_and_logical_and_masked_icmp_asymmetric(i1 %c, i32 %x) {
 
 define i1 @bitwise_and_logical_and_masked_icmp_allzeros(i1 %c, i32 %x) {
 ; CHECK-LABEL: @bitwise_and_logical_and_masked_icmp_allzeros(
-; CHECK-NEXT:    [[X_M1:%.*]] = and i32 [[X:%.*]], 8
-; CHECK-NEXT:    [[C1:%.*]] = icmp eq i32 [[X_M1]], 0
-; CHECK-NEXT:    [[AND1:%.*]] = select i1 [[C1]], i1 [[C:%.*]], i1 false
-; CHECK-NEXT:    [[X_M2:%.*]] = and i32 [[X]], 7
-; CHECK-NEXT:    [[C2:%.*]] = icmp eq i32 [[X_M2]], 0
-; CHECK-NEXT:    [[AND2:%.*]] = and i1 [[AND1]], [[C2]]
-; CHECK-NEXT:    ret i1 [[AND2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP2]], i1 [[C:%.*]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP3]]
 ;
   %x.m1 = and i32 %x, 8
   %c1 = icmp eq i32 %x.m1, 0
@@ -2019,13 +2013,11 @@ define i1 @bitwise_and_logical_and_masked_icmp_allzeros(i1 %c, i32 %x) {
 
 define i1 @bitwise_and_logical_and_masked_icmp_allzeros_poison1(i1 %c, i32 %x, i32 %y) {
 ; CHECK-LABEL: @bitwise_and_logical_and_masked_icmp_allzeros_poison1(
-; CHECK-NEXT:    [[X_M1:%.*]] = and i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[C1:%.*]] = icmp eq i32 [[X_M1]], 0
-; CHECK-NEXT:    [[AND1:%.*]] = select i1 [[C1]], i1 [[C:%.*]], i1 false
-; CHECK-NEXT:    [[X_M2:%.*]] = and i32 [[X]], 7
-; CHECK-NEXT:    [[C2:%.*]] = icmp eq i32 [[X_M2]], 0
-; CHECK-NEXT:    [[AND2:%.*]] = and i1 [[AND1]], [[C2]]
-; CHECK-NEXT:    ret i1 [[AND2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[Y:%.*]], 7
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP2]], 0
+; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP3]], i1 [[C:%.*]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %x.m1 = and i32 %x, %y
   %c1 = icmp eq i32 %x.m1, 0
@@ -2057,13 +2049,10 @@ define i1 @bitwise_and_logical_and_masked_icmp_allzeros_poison2(i1 %c, i32 %x, i
 
 define i1 @bitwise_and_logical_and_masked_icmp_allones(i1 %c, i32 %x) {
 ; CHECK-LABEL: @bitwise_and_logical_and_masked_icmp_allones(
-; CHECK-NEXT:    [[X_M1:%.*]] = and i32 [[X:%.*]], 8
-; CHECK-NEXT:    [[C1:%.*]] = icmp ne i32 [[X_M1]], 0
-; CHECK-NEXT:    [[AND1:%.*]] = select i1 [[C1]], i1 [[C:%.*]], i1 false
-; CHECK-NEXT:    [[X_M2:%.*]] = and i32 [[X]], 7
-; CHECK-NEXT:    [[C2:%.*]] = icmp eq i32 [[X_M2]], 7
-; CHECK-NEXT:    [[AND2:%.*]] = and i1 [[AND1]], [[C2]]
-; CHECK-NEXT:    ret i1 [[AND2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 15
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], 15
+; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP2]], i1 [[C:%.*]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP3]]
 ;
   %x.m1 = and i32 %x, 8
   %c1 = icmp eq i32 %x.m1, 8
@@ -2076,13 +2065,11 @@ define i1 @bitwise_and_logical_and_masked_icmp_allones(i1 %c, i32 %x) {
 
 define i1 @bitwise_and_logical_and_masked_icmp_allones_poison1(i1 %c, i32 %x, i32 %y) {
 ; CHECK-LABEL: @bitwise_and_logical_and_masked_icmp_allones_poison1(
-; CHECK-NEXT:    [[X_M1:%.*]] = and i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[C1:%.*]] = icmp eq i32 [[X_M1]], [[Y]]
-; CHECK-NEXT:    [[AND1:%.*]] = select i1 [[C1]], i1 [[C:%.*]], i1 false
-; CHECK-NEXT:    [[X_M2:%.*]] = and i32 [[X]], 7
-; CHECK-NEXT:    [[C2:%.*]] = icmp eq i32 [[X_M2]], 7
-; CHECK-NEXT:    [[AND2:%.*]] = and i1 [[AND1]], [[C2]]
-; CHECK-NEXT:    ret i1 [[AND2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = or i32 [[Y:%.*]], 7
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i32 [[TMP2]], [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP3]], i1 [[C:%.*]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP4]]
 ;
   %x.m1 = and i32 %x, %y
   %c1 = icmp eq i32 %x.m1, %y