[InstCombine] Simplify cttz/ctlz + icmp eq/ne into mask check

[nikic]

Summary

Related to https://bugs.llvm.org/show_bug.cgi?id=28668.

Checking whether a number has a certain number of trailing / leading zeros means checking whether it is of the form 0001XXXX or XXXX1000, which can be done with an and+icmp.

This is the first step, handling only eq/ne predicates, as that's what the existing code does. I'd like to extend this to handle all predicate types in the next step, which would resolve https://github.com/rust-lang/rust/issues/43024.

[spatel]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2796

We have to check hasOneUse() on the intrinsic, or we'll end up with extra instructions for an example like this:

define i1 @ctlz_eq_other_i32(i32 %x, i32* %p) {
  %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
  store i32 %lz, i32* %p
  %cmp = icmp eq i32 %lz, 24
  ret i1 %cmp
}

In cases where the 'and' gets removed, you don't need the one-use restriction, but I'm not sure if it's worth trying to differentiate those from the general case.

[nikic]

Add single use check and tests.

[nikic]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2796

Thanks, I've added the hasOneUse() check and a couple tests. And yeah, I don't think it's worth handling those cases separately.

[spatel]

LGTM

lib/Transforms/InstCombine/InstCombineCompares.cpp

 Instruction *InstCombiner::foldICmpIntrinsicWithConstant(ICmpInst &Cmp,
                                                          const APInt &C) {
   IntrinsicInst *II = dyn_cast<IntrinsicInst>(Cmp.getOperand(0));
   if (!II || !Cmp.isEquality())
     return nullptr;
 
   // Handle icmp {eq|ne} <intrinsic>, Constant.
   Type *Ty = II->getType();
+  unsigned BitWidth = C.getBitWidth();
   switch (II->getIntrinsicID()) {
   case Intrinsic::bswap:
     Worklist.Add(II);
     Cmp.setOperand(0, II->getArgOperand(0));
     Cmp.setOperand(1, ConstantInt::get(Ty, C.byteSwap()));
     return &Cmp;
 
   case Intrinsic::ctlz:
-  case Intrinsic::cttz:
+  case Intrinsic::cttz: {
     // ctz(A) == bitwidth(A)  ->  A == 0 and likewise for !=
-    if (C == C.getBitWidth()) {
+    if (C == BitWidth) {
       Worklist.Add(II);
       Cmp.setOperand(0, II->getArgOperand(0));
       Cmp.setOperand(1, ConstantInt::getNullValue(Ty));
       return &Cmp;
     }
+
+    // ctz(A) == C -> A & Mask1 == Mask2, where Mask2 only has bit C set
+    // and Mask1 has bits 0..C+1 set. Similar for ctl, but for high bits.
+    // Limit to one use to ensure we don't increase instruction count.
+    unsigned Num = C.getLimitedValue(BitWidth);
+    if (Num != BitWidth && II->hasOneUse()) {
+      bool IsTrailing = II->getIntrinsicID() == Intrinsic::cttz;
+      APInt Mask1 = IsTrailing ? APInt::getLowBitsSet(BitWidth, Num + 1)
+                               : APInt::getHighBitsSet(BitWidth, Num + 1);
+      APInt Mask2 = IsTrailing
+        ? APInt::getOneBitSet(BitWidth, Num)
+        : APInt::getOneBitSet(BitWidth, BitWidth - Num - 1);
+      Cmp.setOperand(0, Builder.CreateAnd(II->getArgOperand(0), Mask1));
+      Cmp.setOperand(1, ConstantInt::get(Ty, Mask2));
+      Worklist.Add(II);
+      return &Cmp;
+    }
     break;
+  }
 
   case Intrinsic::ctpop: {
     // popcount(A) == 0  ->  A == 0 and likewise for !=
     // popcount(A) == bitwidth(A)  ->  A == -1 and likewise for !=
     bool IsZero = C.isNullValue();
-    if (IsZero || C == C.getBitWidth()) {
+    if (IsZero || C == BitWidth) {
       Worklist.Add(II);
       Cmp.setOperand(0, II->getArgOperand(0));
       auto *NewOp =
           IsZero ? Constant::getNullValue(Ty) : Constant::getAllOnesValue(Ty);
       Cmp.setOperand(1, NewOp);
       return &Cmp;
     }
     break;

test/Transforms/InstCombine/cmp-intrinsic.ll

 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %cmp = icmp eq i32 %lz, 32
   ret i1 %cmp
 }
 
 define i1 @ctlz_eq_zero_i32(i32 %x) {
 ; CHECK-LABEL: @ctlz_eq_zero_i32(
-; CHECK-NEXT:    [[LZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range !0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LZ]], 0
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i32 [[X:%.*]], 0
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %cmp = icmp eq i32 %lz, 0
   ret i1 %cmp
 }
 
 define <2 x i1> @ctlz_ne_zero_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @ctlz_ne_zero_v2i32(
-; CHECK-NEXT:    [[X:%.*]] = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> [[A:%.*]], i1 false)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[X]], zeroinitializer
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt <2 x i32> [[A:%.*]], <i32 -1, i32 -1>
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, zeroinitializer
   ret <2 x i1> %cmp
 }
 
 define i1 @ctlz_eq_bw_minus_1_i32(i32 %x) {
 ; CHECK-LABEL: @ctlz_eq_bw_minus_1_i32(
-; CHECK-NEXT:    [[LZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range !0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LZ]], 31
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], 1
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %cmp = icmp eq i32 %lz, 31
   ret i1 %cmp
 }
 
 define <2 x i1> @ctlz_ne_bw_minus_1_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @ctlz_ne_bw_minus_1_v2i32(
-; CHECK-NEXT:    [[X:%.*]] = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> [[A:%.*]], i1 false)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[X]], <i32 31, i32 31>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[A:%.*]], <i32 1, i32 1>
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, <i32 31, i32 31>
   ret <2 x i1> %cmp
 }
 
 define i1 @ctlz_eq_other_i32(i32 %x) {
 ; CHECK-LABEL: @ctlz_eq_other_i32(
-; CHECK-NEXT:    [[LZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range !0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LZ]], 24
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], -128
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[TMP1]], 128
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %cmp = icmp eq i32 %lz, 24
   ret i1 %cmp
 }
 
 define <2 x i1> @ctlz_ne_other_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @ctlz_ne_other_v2i32(
-; CHECK-NEXT:    [[X:%.*]] = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> [[A:%.*]], i1 false)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[X]], <i32 24, i32 24>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[A:%.*]], <i32 -128, i32 -128>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[TMP1]], <i32 128, i32 128>
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, <i32 24, i32 24>
   ret <2 x i1> %cmp
 }
 
+define i1 @ctlz_eq_other_i32_multiuse(i32 %x, i32* %p) {
+; CHECK-LABEL: @ctlz_eq_other_i32_multiuse(
+; CHECK-NEXT:    [[LZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range !0
+; CHECK-NEXT:    store i32 [[LZ]], i32* [[P:%.*]], align 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LZ]], 24
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
+  store i32 %lz, i32* %p
+  %cmp = icmp eq i32 %lz, 24
+  ret i1 %cmp
+}
+
 define <2 x i1> @ctlz_ne_bitwidth_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @ctlz_ne_bitwidth_v2i32(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[A:%.*]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.ctlz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, <i32 32, i32 32>
   ret <2 x i1> %cmp

 ;
   %x = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp eq <2 x i32> %x, <i32 32, i32 32>
   ret <2 x i1> %cmp
 }
 
 define i1 @cttz_eq_zero_i33(i33 %x) {
 ; CHECK-LABEL: @cttz_eq_zero_i33(
-; CHECK-NEXT:    [[TZ:%.*]] = tail call i33 @llvm.cttz.i33(i33 [[X:%.*]], i1 false), !range !1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i33 [[TZ]], 0
+; CHECK-NEXT:    [[TMP1:%.*]] = and i33 [[X:%.*]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i33 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %tz = tail call i33 @llvm.cttz.i33(i33 %x, i1 false)
   %cmp = icmp eq i33 %tz, 0
   ret i1 %cmp
 }
 
 define <2 x i1> @cttz_ne_zero_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @cttz_ne_zero_v2i32(
-; CHECK-NEXT:    [[X:%.*]] = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> [[A:%.*]], i1 false)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[X]], zeroinitializer
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[A:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i32> [[TMP1]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, zeroinitializer
   ret <2 x i1> %cmp
 }
 
 define i1 @cttz_eq_bw_minus_1_i33(i33 %x) {
 ; CHECK-LABEL: @cttz_eq_bw_minus_1_i33(
-; CHECK-NEXT:    [[TZ:%.*]] = tail call i33 @llvm.cttz.i33(i33 [[X:%.*]], i1 false), !range !1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i33 [[TZ]], 32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i33 [[X:%.*]], -4294967296
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %tz = tail call i33 @llvm.cttz.i33(i33 %x, i1 false)
   %cmp = icmp eq i33 %tz, 32
   ret i1 %cmp
 }
 
 define <2 x i1> @cttz_ne_bw_minus_1_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @cttz_ne_bw_minus_1_v2i32(
-; CHECK-NEXT:    [[X:%.*]] = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> [[A:%.*]], i1 false)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[X]], <i32 31, i32 31>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[A:%.*]], <i32 -2147483648, i32 -2147483648>
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, <i32 31, i32 31>
   ret <2 x i1> %cmp
 }
 
 define i1 @cttz_eq_other_i33(i33 %x) {
 ; CHECK-LABEL: @cttz_eq_other_i33(
-; CHECK-NEXT:    [[TZ:%.*]] = tail call i33 @llvm.cttz.i33(i33 [[X:%.*]], i1 false), !range !1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i33 [[TZ]], 4
+; CHECK-NEXT:    [[TMP1:%.*]] = and i33 [[X:%.*]], 31
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i33 [[TMP1]], 16
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %tz = tail call i33 @llvm.cttz.i33(i33 %x, i1 false)
   %cmp = icmp eq i33 %tz, 4
   ret i1 %cmp
 }
 
 define <2 x i1> @cttz_ne_other_v2i32(<2 x i32> %a) {
 ; CHECK-LABEL: @cttz_ne_other_v2i32(
-; CHECK-NEXT:    [[X:%.*]] = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> [[A:%.*]], i1 false)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[X]], <i32 4, i32 4>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[A:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne <2 x i32> [[TMP1]], <i32 16, i32 16>
 ; CHECK-NEXT:    ret <2 x i1> [[CMP]]
 ;
   %x = tail call <2 x i32> @llvm.cttz.v2i32(<2 x i32> %a, i1 false)
   %cmp = icmp ne <2 x i32> %x, <i32 4, i32 4>
   ret <2 x i1> %cmp
 }
 
+define i1 @cttz_eq_other_i33_multiuse(i33 %x, i33* %p) {
+; CHECK-LABEL: @cttz_eq_other_i33_multiuse(
+; CHECK-NEXT:    [[LZ:%.*]] = tail call i33 @llvm.cttz.i33(i33 [[X:%.*]], i1 false), !range !1
+; CHECK-NEXT:    store i33 [[LZ]], i33* [[P:%.*]], align 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i33 [[LZ]], 4
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %lz = tail call i33 @llvm.cttz.i33(i33 %x, i1 false)
+  store i33 %lz, i33* %p
+  %cmp = icmp eq i33 %lz, 4
+  ret i1 %cmp
+}
+
 define i1 @ctpop_eq_zero_i11(i11 %x) {
 ; CHECK-LABEL: @ctpop_eq_zero_i11(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i11 [[X:%.*]], 0
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %pop = tail call i11 @llvm.ctpop.i11(i11 %x)
   %cmp = icmp eq i11 %pop, 0
   ret i1 %cmp

test/Transforms/InstCombine/intrinsics.ll

 ; CHECK-NEXT:    ret i1 false
 ;
   %or = or i32 %arg, 4
   %cnt = call i32 @llvm.cttz.i32(i32 %or, i1 true) nounwind readnone
   %res = icmp eq i32 %cnt, 3
   ret i1 %res
 }
 
-; TODO: The icmp is unnecessary given the known bits of the input.
 define <2 x i1> @cttz_knownbits3_vec(<2 x i32> %arg) {
 ; CHECK-LABEL: @cttz_knownbits3_vec(
-; CHECK-NEXT:    [[OR:%.*]] = or <2 x i32> [[ARG:%.*]], <i32 4, i32 4>
-; CHECK-NEXT:    [[CNT:%.*]] = call <2 x i32> @llvm.cttz.v2i32(<2 x i32> [[OR]], i1 true)
-; CHECK-NEXT:    [[RES:%.*]] = icmp eq <2 x i32> [[CNT]], <i32 3, i32 3>
-; CHECK-NEXT:    ret <2 x i1> [[RES]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %or = or <2 x i32> %arg, <i32 4, i32 4>
   %cnt = call <2 x i32> @llvm.cttz.v2i32(<2 x i32> %or, i1 true) nounwind readnone
   %res = icmp eq <2 x i32> %cnt, <i32 3, i32 3>
   ret <2 x i1> %res
 }
 
 define i8 @ctlz(i8 %a) {

 ; CHECK-NEXT:    ret i1 false
 ;
   %or = or i8 %arg, 32
   %cnt = call i8 @llvm.ctlz.i8(i8 %or, i1 true) nounwind readnone
   %res = icmp eq i8 %cnt, 3
   ret i1 %res
 }
 
-; TODO: The icmp is unnecessary given the known bits of the input.
 define <2 x i1> @ctlz_knownbits3_vec(<2 x i8> %arg) {
 ; CHECK-LABEL: @ctlz_knownbits3_vec(
-; CHECK-NEXT:    [[OR:%.*]] = or <2 x i8> [[ARG:%.*]], <i8 32, i8 32>
-; CHECK-NEXT:    [[CNT:%.*]] = call <2 x i8> @llvm.ctlz.v2i8(<2 x i8> [[OR]], i1 true)
-; CHECK-NEXT:    [[RES:%.*]] = icmp eq <2 x i8> [[CNT]], <i8 3, i8 3>
-; CHECK-NEXT:    ret <2 x i1> [[RES]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %or = or <2 x i8> %arg, <i8 32, i8 32>
   %cnt = call <2 x i8> @llvm.ctlz.v2i8(<2 x i8> %or, i1 true) nounwind readnone
   %res = icmp eq <2 x i8> %cnt, <i8 3, i8 3>
   ret <2 x i1> %res
 }
 
 define i32 @ctlz_undef(i32 %Value) {