[InstCombine] fold icmp of zext bool based on limited range

[spatel]

Summary

Op0 <u (zext i1 X) --> Op0 == 0 && X

https://alive2.llvm.org/ce/z/avQDRY

This is a generalization of 4069cccf3b4ff4a based on the post-commit suggestion.
This also adds the i1 type check and tests that were missing from the earlier attempt; that commit caused several bot fails and was reverted.

[nikic]

This comment has been deleted.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5647

Might make sense to use m_c_ICmp instead? Something along these lines:

if (match(I, m_c_ICmp(Pred, m_Value(X), m_OneUse(m_ZExt(m_Value(Y))))) &&
    Y->getType()->isIntOrIntVectorTy(1) && Pred == ICmpInst::ICMP_ULT)

That would probably make it easier to extend this to the conjugate sext pattern: https://alive2.llvm.org/ce/z/2JTkdF

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

5647

Ah, thanks - I knew there was a sibling, but I had not worked out the pattern. I'll add a TODO for that.

And yes, we can use the commutative matcher. I wasn't sure if we'd miss the transform if both operands are zexted, but it looks like other folds will convert that to something we can match more easily here. I'll add some more tests to confirm.

[spatel]

Patch updated:

1. Use commutative matcher to reduce code.
2. Add TODO for related pattern.
3. Verify fold with more tests.

[nikic]

LG

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombinerImpl::foldICmpUsingKnownBits(ICmpInst &I) {
   if (I.isSigned() &&
       ((Op0Known.Zero.isNegative() && Op1Known.Zero.isNegative()) ||
        (Op0Known.One.isNegative() && Op1Known.One.isNegative())))
     return new ICmpInst(I.getUnsignedPredicate(), Op0, Op1);
 
   return nullptr;
 }
 
+/// If one operand of an icmp is effectively a bool (value range of {0,1}),
+/// then try to reduce patterns based on that limit.
+static Instruction *foldICmpUsingBoolRange(ICmpInst &I,
+                                           InstCombiner::BuilderTy &Builder) {
+  Value *X, *Y;
+  ICmpInst::Predicate Pred;
+
+  // X must be 0 and bool must be true for "ULT":
+  // X <u (zext i1 Y) --> (X == 0) && Y
+  if (match(&I, m_c_ICmp(Pred, m_Value(X), m_OneUse(m_ZExt(m_Value(Y))))) &&
+      Y->getType()->isIntOrIntVectorTy(1) && Pred == ICmpInst::ICMP_ULT)
+    return BinaryOperator::CreateAnd(Builder.CreateIsNull(X), Y);
+
+  // TODO: Handle the related pattern with UGE/sext.
+
+  return nullptr;
+}
+
 llvm::Optional<std::pair<CmpInst::Predicate, Constant *>>
 InstCombiner::getFlippedStrictnessPredicateAndConstant(CmpInst::Predicate Pred,
                                                        Constant *C) {
   assert(ICmpInst::isRelational(Pred) && ICmpInst::isIntPredicate(Pred) &&
          "Only for relational integer predicates.");
 
   Type *Type = C->getType();
   bool IsSigned = ICmpInst::isSigned(Pred);
if (Instruction *Res = canonicalizeICmpPredicate(I))
     return Res;
 
   if (Instruction *Res = foldICmpWithConstant(I))
     return Res;
 
   if (Instruction *Res = foldICmpWithDominatingICmp(I))
     return Res;
 
+  if (Instruction *Res = foldICmpUsingBoolRange(I, Builder))
+    return Res;
+
   if (Instruction *Res = foldICmpUsingKnownBits(I))
     return Res;
 
   // Test if the ICmpInst instruction is used exclusively by a select as
   // part of a minimum or maximum operation. If so, refrain from doing
   // any other folding. This helps out other analyses which understand
   // non-obfuscated minimum and maximum idioms, such as ScalarEvolution
   // and CodeGen. And in this case, at least one of the comparison

llvm/test/Transforms/InstCombine/icmp-range.ll

;
   %val1 = load i32, i32* %arg1, !range !0
   %val2 = load i32, i32* %arg2, !range !6
   %rval = icmp ugt i32 %val2, %val1
   ret i1 %rval
 }
 
 define i1 @ugt_zext(i1 %b, i8 %x) {
 ; CHECK-LABEL: @ugt_zext(
-; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[B:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[Z]], [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[X:%.*]], 0
+; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP1]], [[B:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i1 %b to i8
   %r = icmp ugt i8 %z, %x
   ret i1 %r
 }
 
 define <2 x i1> @ult_zext(<2 x i1> %b, <2 x i8> %p) {
 ; CHECK-LABEL: @ult_zext(
 ; CHECK-NEXT:    [[X:%.*]] = mul <2 x i8> [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i1> [[B:%.*]] to <2 x i8>
-; CHECK-NEXT:    [[R:%.*]] = icmp ult <2 x i8> [[X]], [[Z]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <2 x i8> [[X]], zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = and <2 x i1> [[TMP1]], [[B:%.*]]
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %x = mul <2 x i8> %p, %p ; thwart complexity-based canonicalization
   %z = zext <2 x i1> %b to <2 x i8>
   %r = icmp ult <2 x i8> %x, %z
   ret <2 x i1> %r
 }
 
+; negative test - need ult/ugt
+
 define i1 @uge_zext(i1 %b, i8 %x) {
 ; CHECK-LABEL: @uge_zext(
 ; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[B:%.*]] to i8
 ; CHECK-NEXT:    [[R:%.*]] = icmp uge i8 [[Z]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i1 %b to i8
   %r = icmp uge i8 %z, %x
   ret i1 %r
 }
 
+; negative test - need ult/ugt
+
 define i1 @ule_zext(i1 %b, i8 %p) {
 ; CHECK-LABEL: @ule_zext(
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[P:%.*]], [[P]]
 ; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[B:%.*]] to i8
 ; CHECK-NEXT:    [[R:%.*]] = icmp ule i8 [[X]], [[Z]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %x = mul i8 %p, %p ; thwart complexity-based canonicalization
   %z = zext i1 %b to i8
   %r = icmp ule i8 %x, %z
   ret i1 %r
 }
 
+; negative test - extra use
+
 define i1 @ugt_zext_use(i1 %b, i8 %x) {
 ; CHECK-LABEL: @ugt_zext_use(
 ; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[B:%.*]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[Z]])
 ; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[Z]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i1 %b to i8
   call void @use(i8 %z)
   %r = icmp ugt i8 %z, %x
   ret i1 %r
 }
 
+; negative test - must be zext of i1
+
 define i1 @ult_zext_not_i1(i2 %b, i8 %x) {
 ; CHECK-LABEL: @ult_zext_not_i1(
 ; CHECK-NEXT:    [[Z:%.*]] = zext i2 [[B:%.*]] to i8
 ; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[Z]], [[X:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i2 %b to i8
   %r = icmp ult i8 %x, %z
   ret i1 %r
 }
 
+; sub is eliminated
+
 define i1 @sub_ult_zext(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @sub_ult_zext(
-; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[B:%.*]] to i8
-; CHECK-NEXT:    [[S:%.*]] = sub i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[S]], [[Z]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP1]], [[B:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i1 %b to i8
   %s = sub i8 %x, %y
   %r = icmp ult i8 %s, %z
   ret i1 %r
 }
 
 define i1 @zext_ult_zext(i1 %b, i8 %p) {
 ; CHECK-LABEL: @zext_ult_zext(
 ; CHECK-NEXT:    [[X:%.*]] = mul i8 [[P:%.*]], [[P]]
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[B:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[X]], [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[X]], 0
+; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP1]], [[B:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %x = mul i8 %p, %p ; thwart complexity-based canonicalization
   %z = zext i1 %b to i16
   %zx = zext i8 %x to i16
   %r = icmp ult i16 %zx, %z
   ret i1 %r
 }
 
+; match and fold even if both sides are zexts (from different source types)
+
 define i1 @zext_ugt_zext(i1 %b, i4 %x) {
 ; CHECK-LABEL: @zext_ugt_zext(
 ; CHECK-NEXT:    [[ZX:%.*]] = zext i4 [[X:%.*]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[ZX]])
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[B:%.*]] to i4
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt i4 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i4 [[X]], 0
+; CHECK-NEXT:    [[R:%.*]] = and i1 [[TMP1]], [[B:%.*]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i1 %b to i8
   %zx = zext i4 %x to i8
   call void @use(i8 %zx)
   %r = icmp ugt i8 %z, %zx
   ret i1 %r
 }
 
+; negative test - must be zext of i1
+
 define i1 @sub_ult_zext_not_i1(i2 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @sub_ult_zext_not_i1(
 ; CHECK-NEXT:    [[Z:%.*]] = zext i2 [[B:%.*]] to i8
 ; CHECK-NEXT:    [[S:%.*]] = sub i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[S]], [[Z]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i2 %b to i8
   %s = sub i8 %x, %y
   %r = icmp ult i8 %s, %z
   ret i1 %r
 }
 
+; negative test - extra use (but we could try harder to fold this)
+
 define i1 @sub_ult_zext_use1(i1 %b, i8 %x, i8 %y) {
 ; CHECK-LABEL: @sub_ult_zext_use1(
 ; CHECK-NEXT:    [[Z:%.*]] = zext i1 [[B:%.*]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[Z]])
 ; CHECK-NEXT:    [[S:%.*]] = sub i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[S]], [[Z]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %z = zext i1 %b to i8
   call void @use(i8 %z)
   %s = sub i8 %x, %y
   %r = icmp ult i8 %s, %z
   ret i1 %r
 }
 
 define <2 x i1> @zext_ugt_sub_use2(<2 x i1> %b, <2 x i8> %x, <2 x i8> %y) {
 ; CHECK-LABEL: @zext_ugt_sub_use2(
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i1> [[B:%.*]] to <2 x i8>
 ; CHECK-NEXT:    [[S:%.*]] = sub <2 x i8> [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    call void @use_vec(<2 x i8> [[S]])
-; CHECK-NEXT:    [[R:%.*]] = icmp ult <2 x i8> [[S]], [[Z]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <2 x i8> [[X]], [[Y]]
+; CHECK-NEXT:    [[R:%.*]] = and <2 x i1> [[TMP1]], [[B:%.*]]
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %z = zext <2 x i1> %b to <2 x i8>
   %s = sub <2 x i8> %x, %y
   call void @use_vec(<2 x i8> %s)
   %r = icmp ugt <2 x i8> %z, %s
   ret <2 x i1> %r
 }