[InstCombine] Improve fold icmp of truncated left shift

[Chenbing.Zheng]

Summary

Original version only solve shifted constant 1, this patch expend it to PowerOf2

(trunc (2**Z << Y) to iN) == 0 -> Y u>  N - Z + 1
(trunc (2**Z << Y) to iN) != 0 -> Y u<= N - Z + 1

https://alive2.llvm.org/ce/z/zaQFeR
and

(trunc (2**Z << Y) to iN) == 2**C --> Y == C - Z
(trunc (2**Z << Y) to iN) != 2**C --> Y != C - Z

[spatel]

The TODO comment was actually suggesting that we handle something like this:
https://alive2.llvm.org/ce/z/dHssjh
...but this patch won't match that? If that's correct, please leave the TODO comment there.

Please add that plus any other baseline tests, so we are just showing the diffs.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1559

I prefer to give a constant value some name with "C" to make it obvious that we matched a constant, so C1 or ShiftC or something like that.

1560–1561

Move this comment down, so it is clear that it applies only when C == 0 - just above the C.isZero() line.
Change "Z" in the comment to match the variable name if it is changed.

[Chenbing.Zheng]

address comments

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1561

That looks correct, but I'd prefer that this comment matches what the code and the alive proof actually do, so I'd change it to something like this:

// (trunc (2**C1 << Y) to iN) == 0 -> Y u> ctlz(trunc(C1))

1569–1570

Add an Alive2 proof for this pattern to the patch description?

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

409

Remove TODO.

[RKSimon]

@Chenbing.Zheng Are you still looking at this?

[Chenbing.Zheng]

I'm going to abandon this patch. It seems that there is a more complete solution here.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombinerImpl::foldICmpTruncConstant(ICmpInst &Cmp,
 
   Type *SrcTy = X->getType();
   unsigned DstBits = Trunc->getType()->getScalarSizeInBits(),
            SrcBits = SrcTy->getScalarSizeInBits();
 
   // TODO: Handle any shifted constant by subtracting trailing zeros.
   // TODO: Handle non-equality predicates.
   Value *Y;
-  if (Cmp.isEquality() && match(X, m_Shl(m_One(), m_Value(Y)))) {
-    // (trunc (1 << Y) to iN) == 0 --> Y u>= N
-    // (trunc (1 << Y) to iN) != 0 --> Y u<  N
+  const APInt *C1;
+  if (Cmp.isEquality() && match(X, m_Shl(m_APIntAllowUndef(C1), m_Value(Y)))) {
+    if (C1->isPowerOf2()) {
+      // (trunc (2**C1 << Y) to iN) == 0 -> Y u>  N - C1 + 1
+      // (trunc (2**C1 << Y) to iN) != 0 -> Y u<= N - C1 + 1
       if (C.isZero()) {
-      auto NewPred = (Pred == Cmp.ICMP_EQ) ? Cmp.ICMP_UGE : Cmp.ICMP_ULT;
-      return new ICmpInst(NewPred, Y, ConstantInt::get(SrcTy, DstBits));
+        auto NewPred = (Pred == Cmp.ICMP_EQ) ? Cmp.ICMP_UGT : Cmp.ICMP_ULE;
+        return new ICmpInst(
+            NewPred, Y,
+            ConstantInt::get(SrcTy, C1->trunc(DstBits).countLeadingZeros()));
+      }
+      // (trunc (2**C1 << Y) to iN) == 2**C --> Y == C - C1
+      // (trunc (2**C1 << Y) to iN) != 2**C --> Y != C - C1
+      if (C.isPowerOf2() && C.uge(*C1))
+        return new ICmpInst(
+            Pred, Y, ConstantInt::get(SrcTy, C.logBase2() - C1->logBase2()));
     }
-    // (trunc (1 << Y) to iN) == 2**C --> Y == C
-    // (trunc (1 << Y) to iN) != 2**C --> Y != C
-    if (C.isPowerOf2())
-      return new ICmpInst(Pred, Y, ConstantInt::get(SrcTy, C.logBase2()));
   }
 
   if (Cmp.isEquality() && Trunc->hasOneUse()) {
     // Canonicalize to a mask and wider compare if the wide type is suitable:
     // (trunc X to i8) == C --> (X & 0xff) == (zext C)
     if (!SrcTy->isVectorTy() && shouldChangeType(DstBits, SrcBits)) {
       Constant *Mask =
           ConstantInt::get(SrcTy, APInt::getLowBitsSet(SrcBits, DstBits));

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

 ;
   %shl = shl i37 1, %a
   %t = trunc i37 %shl to i8
   call void @use(i8 %t)
   %r = icmp ne i8 %t, 0
   ret i1 %r
 }
 
 ; TODO: A > 4
 
 define i1 @shl2_trunc_eq0(i9 %a) {
 ; CHECK-LABEL: @shl2_trunc_eq0(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i9 2, [[A:%.*]]
-; CHECK-NEXT:    [[T:%.*]] = trunc i9 [[SHL]] to i6
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i6 [[T]], 0
+; CHECK-NEXT:    [[R:%.*]] = icmp ugt i9 [[A:%.*]], 4
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i9 2, %a
   %t = trunc i9 %shl to i6
   %r = icmp eq i6 %t, 0
   ret i1 %r
 }
 
 define i1 @shl2_trunc_ne0(i9 %a) {
 ; CHECK-LABEL: @shl2_trunc_ne0(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i9 2, [[A:%.*]]
-; CHECK-NEXT:    [[T:%.*]] = trunc i9 [[SHL]] to i6
-; CHECK-NEXT:    [[R:%.*]] = icmp ne i6 [[T]], 0
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i9 [[A:%.*]], 5
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i9 2, %a
   %t = trunc i9 %shl to i6
   %r = icmp ne i6 %t, 0
   ret i1 %r
 }
 
 define i1 @shl3_trunc_eq0(i9 %a) {
 ; CHECK-LABEL: @shl3_trunc_eq0(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i9 3, [[A:%.*]]
 ; CHECK-NEXT:    [[T:%.*]] = trunc i9 [[SHL]] to i6
 ; CHECK-NEXT:    [[R:%.*]] = icmp eq i6 [[T]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i9 3, %a
   %t = trunc i9 %shl to i6
   %r = icmp eq i6 %t, 0
   ret i1 %r
 }
 
 define <2 x i1> @shl4_trunc_ne0(<2 x i8> %a) {
 ; CHECK-LABEL: @shl4_trunc_ne0(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i8> <i8 4, i8 poison>, [[A:%.*]]
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i8> [[SHL]] to <2 x i5>
-; CHECK-NEXT:    [[R:%.*]] = icmp ne <2 x i5> [[T]], zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = icmp ult <2 x i8> [[A:%.*]], <i8 3, i8 3>
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %shl = shl <2 x i8> <i8 4, i8 poison>, %a
   %t = trunc <2 x i8> %shl to <2 x i5>
   %r = icmp ne <2 x i5> %t, zeroinitializer
   ret <2 x i1> %r
 }
 
;
   %shl = shl i8 1, %a
   call void @use(i8 %shl)
   %t = trunc i8 %shl to i6
   %r = icmp ne i6 %t, 32
   ret i1 %r
 }
 
 define i1 @shl2_trunc_eq8_i32(i32 %a) {
-; DL64-LABEL: @shl2_trunc_eq8_i32(
-; DL64-NEXT:    [[SHL:%.*]] = shl i32 2, [[A:%.*]]
-; DL64-NEXT:    [[TMP1:%.*]] = and i32 [[SHL]], 65534
-; DL64-NEXT:    [[R:%.*]] = icmp eq i32 [[TMP1]], 8
-; DL64-NEXT:    ret i1 [[R]]
-;
-; DL8-LABEL: @shl2_trunc_eq8_i32(
-; DL8-NEXT:    [[SHL:%.*]] = shl i32 2, [[A:%.*]]
-; DL8-NEXT:    [[T:%.*]] = trunc i32 [[SHL]] to i16
-; DL8-NEXT:    [[R:%.*]] = icmp eq i16 [[T]], 8
-; DL8-NEXT:    ret i1 [[R]]
+; CHECK-LABEL: @shl2_trunc_eq8_i32(
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[A:%.*]], 2
+; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i32 2, %a
   %t = trunc i32 %shl to i16
   %r = icmp eq i16 %t, 8
   ret i1 %r
 }
 
 define i1 @shl2_trunc_ne8_i32(i32 %a) {
-; DL64-LABEL: @shl2_trunc_ne8_i32(
-; DL64-NEXT:    [[SHL:%.*]] = shl i32 2, [[A:%.*]]
-; DL64-NEXT:    [[TMP1:%.*]] = and i32 [[SHL]], 65534
-; DL64-NEXT:    [[R:%.*]] = icmp ne i32 [[TMP1]], 8
-; DL64-NEXT:    ret i1 [[R]]
-;
-; DL8-LABEL: @shl2_trunc_ne8_i32(
-; DL8-NEXT:    [[SHL:%.*]] = shl i32 2, [[A:%.*]]
-; DL8-NEXT:    [[T:%.*]] = trunc i32 [[SHL]] to i16
-; DL8-NEXT:    [[R:%.*]] = icmp ne i16 [[T]], 8
-; DL8-NEXT:    ret i1 [[R]]
+; CHECK-LABEL: @shl2_trunc_ne8_i32(
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i32 [[A:%.*]], 2
+; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i32 2, %a
   %t = trunc i32 %shl to i16
   %r = icmp ne i16 %t, 8
   ret i1 %r
 }
 
 define i1 @shl1_trunc_sgt4(i32 %a) {