[InstCombine] Support logical ops in foldAndOrOfICmpEqZeroAndICmp()

If the and/or is logical and one of the operands only occurs on the RHS, we need to freeze it: https://alive2.llvm.org/ce/z/vuMuE_
apple · Dec 13, 2022 · de59d22 · de59d22
1 parent 004a70a
commit de59d22
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 19 deletions.
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -2584,8 +2584,9 @@ Value *InstCombinerImpl::matchSelectFromAndOr(Value *A, Value *C, Value *B,
 
 // (icmp eq X, 0) | (icmp ult Other, X) -> (icmp ule Other, X-1)
 // (icmp ne X, 0) & (icmp uge Other, X) -> (icmp ugt Other, X-1)
-Value *foldAndOrOfICmpEqZeroAndICmp(ICmpInst *LHS, ICmpInst *RHS, bool IsAnd,
-                                    IRBuilderBase &Builder) {
+static Value *foldAndOrOfICmpEqZeroAndICmp(ICmpInst *LHS, ICmpInst *RHS,
+                                           bool IsAnd, bool IsLogical,
+                                           IRBuilderBase &Builder) {
   ICmpInst::Predicate LPred =
       IsAnd ? LHS->getInversePredicate() : LHS->getPredicate();
   ICmpInst::Predicate RPred =
@@ -2604,6 +2605,8 @@ Value *foldAndOrOfICmpEqZeroAndICmp(ICmpInst *LHS, ICmpInst *RHS, bool IsAnd,
   else
     return nullptr;
 
+  if (IsLogical)
+    Other = Builder.CreateFreeze(Other);
   return Builder.CreateICmp(
       IsAnd ? ICmpInst::ICMP_ULT : ICmpInst::ICMP_UGE,
       Builder.CreateAdd(LHS0, Constant::getAllOnesValue(LHS0->getType())),
@@ -2652,14 +2655,14 @@ Value *InstCombinerImpl::foldAndOrOfICmps(ICmpInst *LHS, ICmpInst *RHS,
   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.
-  if (!IsLogical) {
-    if (Value *V = foldAndOrOfICmpEqZeroAndICmp(LHS, RHS, IsAnd, Builder))
-      return V;
-    if (Value *V = foldAndOrOfICmpEqZeroAndICmp(RHS, LHS, IsAnd, Builder))
-      return V;
-  }
+  if (Value *V =
+          foldAndOrOfICmpEqZeroAndICmp(LHS, RHS, IsAnd, IsLogical, Builder))
+    return V;
+  // We can treat logical like bitwise here, because both operands are used on
+  // the LHS, and as such poison from both will propagate.
+  if (Value *V = foldAndOrOfICmpEqZeroAndICmp(RHS, LHS, IsAnd,
+                                              /*IsLogical*/ false, Builder))
+    return V;
 
   // TODO: Verify whether this is safe for logical and/or.
   if (!IsLogical) {

diff --git a/llvm/test/Transforms/InstCombine/icmp.ll b/llvm/test/Transforms/InstCombine/icmp.ll
@@ -2172,9 +2172,9 @@ define i1 @or_icmp_eq_B_0_icmp_ult_A_B(i64 %a, i64 %b) {
 
 define i1 @or_icmp_eq_B_0_icmp_ult_A_B_logical(i64 %a, i64 %b) {
 ; CHECK-LABEL: @or_icmp_eq_B_0_icmp_ult_A_B_logical(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[B:%.*]], 0
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i64 [[A:%.*]], [[B]]
-; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i1 true, i1 [[TMP2]]
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i64 [[A:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[B:%.*]], -1
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp uge i64 [[TMP2]], [[TMP1]]
 ; CHECK-NEXT:    ret i1 [[TMP3]]
 ;
   %1 = icmp eq i64 %b, 0
@@ -2306,9 +2306,9 @@ define i1 @and_icmp_ne_B_0_icmp_uge_A_B_commuted1_logical(i64 %a, i64 %b) {
 
 define i1 @and_icmp_ne_B_0_icmp_uge_A_B_commuted2_logical(i64 %a, i64 %b) {
 ; CHECK-LABEL: @and_icmp_ne_B_0_icmp_uge_A_B_commuted2_logical(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i64 [[B:%.*]], 0
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ule i64 [[B]], [[A:%.*]]
-; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i1 [[TMP2]], i1 false
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i64 [[A:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[B:%.*]], -1
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult i64 [[TMP2]], [[TMP1]]
 ; CHECK-NEXT:    ret i1 [[TMP3]]
 ;
   %1 = icmp ne i64 %b, 0
@@ -2418,9 +2418,9 @@ define i1 @and_icmp_ne_B_0_icmp_uge_A_B_wrong_op2(i64 %a, i64 %b, i64 %c) {
 
 define i1 @and_icmp_ne_B_0_icmp_uge_A_B_logical(i64 %a, i64 %b) {
 ; CHECK-LABEL: @and_icmp_ne_B_0_icmp_uge_A_B_logical(
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i64 [[B:%.*]], 0
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp uge i64 [[A:%.*]], [[B]]
-; CHECK-NEXT:    [[TMP3:%.*]] = select i1 [[TMP1]], i1 [[TMP2]], i1 false
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i64 [[A:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[B:%.*]], -1
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ult i64 [[TMP2]], [[TMP1]]
 ; CHECK-NEXT:    ret i1 [[TMP3]]
 ;
   %1 = icmp ne i64 %b, 0