[InstCombine] Fold icmp of truncated left shift

[spatel]

Summary

Refer to issue #51889. Folds:

• (icmp eq (trunc (shl 1, Y) to i[N]), 0) --> (icmp ugt Y, N-1)
• (icmp ne (trunc (shl 1, Y) to i[N]), 0) --> (icmp ult Y, N)

[hasyimibhar]

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

215 ↗

(On Diff #393327)

Just noticed the output is different for ne with truncate to i1. According to alive, this is correct: https://alive2.llvm.org/ce/z/pWkv74.

[hasyimibhar]

Run git-clang-format

[RKSimon]

Please can you add some vector tests? And also some negative (multi-use) tests?

[spatel]

There are a lot of tests here using different bit-widths, but they don't add much value. As suggested, we'd get more impact by varying types/uses (hint: if you're only creating 1 instruction, then extra uses of the intermediate values are almost always ok).

Also, as noted in the bug report, the compare to 0 is possibly a special-case of a more general pattern. Did you consider handling something like this:
https://alive2.llvm.org/ce/z/6SCoJu ?

Finally, if you don't have commit access yet, you probably have enough commits now to request it:
https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access

[spatel]

Reverse ping @hasyimibhar - do you plan to continue work on this patch, or should someone else commandeer it?

[spatel]

Patch updated:

1. Added/removed/pre-committed tests.
2. Adjusted position of fold.

The diffs show that we can already fold some of these patterns via other transforms, but it is limited based on data type. That seems like an unnecessary restriction for this pattern since we can reduce to a single instruction (no one-use limits either). There are many TODO enhancements that can be done in follow-up patches.

[RKSimon]

LGTM - cheers

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

if (C.isOne() && C.getBitWidth() > 1) {
     if (Pred == ICmpInst::ICMP_SLT && match(X, m_Signum(m_Value(V))))
       return new ICmpInst(ICmpInst::ICMP_SLT, V,
                           ConstantInt::get(V->getType(), 1));
   }
 
   Type *SrcTy = X->getType();
   unsigned DstBits = Trunc->getType()->getScalarSizeInBits(),
            SrcBits = SrcTy->getScalarSizeInBits();
+
+  // (trunc (1 << Y) to iN) == 0 --> Y u>= N
+  // (trunc (1 << Y) to iN) != 0 --> Y u<  N
+  // TODO: Handle any shifted constant by subtracting trailing zeros.
+  // TODO: Handle non-equality predicates.
+  // TODO: Handle compare to power-of-2 (non-zero) constant.
+  Value *Y;
+  if (Cmp.isEquality() && C.isZero() && match(X, m_Shl(m_One(), m_Value(Y)))) {
+    auto NewPred = Pred == Cmp.ICMP_EQ ? Cmp.ICMP_UGE : Cmp.ICMP_ULT;
+    return new ICmpInst(NewPred, Y, ConstantInt::get(SrcTy, DstBits));
+  }
+
   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));
       Value *And = Builder.CreateAnd(X, Mask);
       Constant *WideC = ConstantInt::get(SrcTy, C.zext(SrcBits));

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

 ; DL8-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i64 %x to i10
   %r = icmp eq i10 %t, 42
   ret i1 %r
 }
 
 define i1 @shl1_trunc_eq0(i32 %a) {
-; DL64-LABEL: @shl1_trunc_eq0(
-; DL64-NEXT:    [[R:%.*]] = icmp ugt i32 [[A:%.*]], 15
-; DL64-NEXT:    ret i1 [[R]]
-;
-; DL8-LABEL: @shl1_trunc_eq0(
-; DL8-NEXT:    [[SHL:%.*]] = shl i32 1, [[A:%.*]]
-; DL8-NEXT:    [[T:%.*]] = trunc i32 [[SHL]] to i16
-; DL8-NEXT:    [[R:%.*]] = icmp eq i16 [[T]], 0
-; DL8-NEXT:    ret i1 [[R]]
+; CHECK-LABEL: @shl1_trunc_eq0(
+; CHECK-NEXT:    [[R:%.*]] = icmp ugt i32 [[A:%.*]], 15
+; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i32 1, %a
   %t = trunc i32 %shl to i16
   %r = icmp eq i16 %t, 0
   ret i1 %r
 }
 
 define <2 x i1> @shl1_trunc_ne0(<2 x i8> %a) {
 ; CHECK-LABEL: @shl1_trunc_ne0(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i8> <i8 1, 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 5, i8 5>
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %shl = shl <2 x i8> <i8 1, 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
 }
 
;
   ret i1 %r
 }
 
 define i1 @shl1_trunc_ne0_use2(i37 %a) {
 ; CHECK-LABEL: @shl1_trunc_ne0_use2(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i37 1, [[A:%.*]]
 ; CHECK-NEXT:    [[T:%.*]] = trunc i37 [[SHL]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[T]])
-; CHECK-NEXT:    [[R:%.*]] = icmp ne i8 [[T]], 0
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i37 [[A]], 8
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %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:    ret i1 [[R]]
 ;
   %shl = shl i9 2, %a
   %t = trunc i9 %shl to i6
   %r = icmp eq i6 %t, 0
   ret i1 %r
 }
 
+; TODO: A < 5
+
 define i1 @shl1_trunc_sgt0(i9 %a) {
 ; CHECK-LABEL: @shl1_trunc_sgt0(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i9 1, [[A:%.*]]
 ; CHECK-NEXT:    [[T:%.*]] = trunc i9 [[SHL]] to i6
 ; CHECK-NEXT:    [[R:%.*]] = icmp sgt i6 [[T]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i9 1, %a
   %t = trunc i9 %shl to i6
   %r = icmp sgt i6 %t, 0
   ret i1 %r
 }
 
+; TODO: A == 0
+
 define i1 @shl1_trunc_eq1(i64 %a) {
 ; CHECK-LABEL: @shl1_trunc_eq1(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i64 1, [[A:%.*]]
 ; CHECK-NEXT:    [[T:%.*]] = trunc i64 [[SHL]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    [[R:%.*]] = icmp eq i8 [[T]], 1
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i64 1, %a
   %t = trunc i64 %shl to i8
   call void @use(i8 %t)
   %r = icmp eq i8 %t, 1
   ret i1 %r
 }
 
+; TODO: A != 5
+
 define i1 @shl1_trunc_ne32(i8 %a) {
 ; CHECK-LABEL: @shl1_trunc_ne32(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i8 1, [[A:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[SHL]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[SHL]], 63
 ; CHECK-NEXT:    [[R:%.*]] = icmp ne i8 [[TMP1]], 32
 ; CHECK-NEXT:    ret i1 [[R]]
 ;