[InstCombine] Generalize trunc-shift-icmp fold from (1 << Y) to (Pow2 << Y) (llvm#169163)

Extends the `icmp(trunc(shl))` fold to handle any power of 2 constant as
the shift base, not just 1. This generalizes the following patterns by
adjusting the comparison offsets by `log2(Pow2)`.

```llvm
(trunc (1 << Y) to iN) == 0    --> Y u>= N
(trunc (1 << Y) to iN) != 0    --> Y u<  N
(trunc (1 << Y) to iN) == 2**C --> Y ==  C
(trunc (1 << Y) to iN) != 2**C --> Y !=  C

; to

(trunc (Pow2 << Y) to iN) == 0    --> Y u>= N - log2(Pow2)
(trunc (Pow2 << Y) to iN) != 0    --> Y u<  N - log2(Pow2)
(trunc (Pow2 << Y) to iN) == 2**C --> Y ==  C - log2(Pow2)
(trunc (Pow2 << Y) to iN) != 2**C --> Y !=  C - log2(Pow2)
```

Proof: https://alive2.llvm.org/ce/z/2zwTkp

Author: pedroclobo

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

@@ -1465,20 +1465,24 @@ Instruction *InstCombinerImpl::foldICmpTruncConstant(ICmpInst &Cmp,
                           ConstantInt::get(V->getType(), 1));
   }
 
-  // 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 *Pow2;
+  if (Cmp.isEquality() && match(X, m_Shl(m_Power2(Pow2), m_Value(Y))) &&
+      DstBits > Pow2->logBase2()) {
+    // (trunc (Pow2 << Y) to iN) == 0 --> Y u>= N - log2(Pow2)
+    // (trunc (Pow2 << Y) to iN) != 0 --> Y u<  N - log2(Pow2)
+    // iff N > log2(Pow2)
     if (C.isZero()) {
       auto NewPred = (Pred == Cmp.ICMP_EQ) ? Cmp.ICMP_UGE : Cmp.ICMP_ULT;
-      return new ICmpInst(NewPred, Y, ConstantInt::get(SrcTy, DstBits));
+      return new ICmpInst(NewPred, Y,
+                          ConstantInt::get(SrcTy, DstBits - Pow2->logBase2()));
     }
-    // (trunc (1 << Y) to iN) == 2**C --> Y == C
-    // (trunc (1 << Y) to iN) != 2**C --> Y != C
+    // (trunc (Pow2 << Y) to iN) == 2**C --> Y == C - log2(Pow2)
+    // (trunc (Pow2 << Y) to iN) != 2**C --> Y != C - log2(Pow2)
     if (C.isPowerOf2())
-      return new ICmpInst(Pred, Y, ConstantInt::get(SrcTy, C.logBase2()));
+      return new ICmpInst(
+          Pred, Y, ConstantInt::get(SrcTy, C.logBase2() - Pow2->logBase2()));
   }
 
   if (Cmp.isEquality() && (Trunc->hasOneUse() || Trunc->hasNoUnsignedWrap())) {

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

@@ -407,13 +407,9 @@ define i1 @shl1_trunc_ne0_use2(i37 %a) {
   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
@@ -424,9 +420,7 @@ define i1 @shl2_trunc_eq0(i9 %a) {
 
 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
@@ -450,9 +444,7 @@ define i1 @shl3_trunc_eq0(i9 %a) {
 
 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:%.*]], splat (i8 3)
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %shl = shl <2 x i8> <i8 4, i8 poison>, %a
@@ -461,6 +453,16 @@ define <2 x i1> @shl4_trunc_ne0(<2 x i8> %a) {
   ret <2 x i1> %r
 }
 
+define i1 @shl5_trunc_ne0(i9 %a) {
+; CHECK-LABEL: @shl5_trunc_ne0(
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i9 [[A:%.*]], 4
+; CHECK-NEXT:    ret i1 [[R]]
+;
+  %shl = shl i9 4, %a
+  %t = trunc i9 %shl to i6
+  %r = icmp ne i6 %t, 0
+  ret i1 %r
+}
 
 ; TODO: A < 5
 
@@ -507,17 +509,9 @@ define i1 @shl1_trunc_ne32(i8 %a) {
 }
 
 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
@@ -526,24 +520,36 @@ define i1 @shl2_trunc_eq8_i32(i32 %a) {
 }
 
 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 @neg_shl2_trunc_eq0_i8(i8 %a) {
+; CHECK-LABEL: @neg_shl2_trunc_eq0_i8(
+; CHECK-NEXT:    ret i1 true
+;
+  %shl = shl i8 128, %a
+  %t = trunc i8 %shl to i6
+  %r = icmp eq i6 %t, 0
+  ret i1 %r
+}
+
+define i1 @neg_shl2_trunc_ne0_i8(i8 %a) {
+; CHECK-LABEL: @neg_shl2_trunc_ne0_i8(
+; CHECK-NEXT:    ret i1 false
+;
+  %shl = shl i8 128, %a
+  %t = trunc i8 %shl to i6
+  %r = icmp ne i6 %t, 0
+  ret i1 %r
+}
+
 define i1 @shl1_trunc_sgt4(i32 %a) {
 ; CHECK-LABEL: @shl1_trunc_sgt4(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i32 1, [[A:%.*]]