[InstCombine] Look through truncate to fold icmp with intrinsics

[bcl5980]

Summary

The output of intrinsic functions like ctpop, cttz, ctlz have limited range from 0 to bitwidth. So if the truncate destination type can hold the source bitwidth size, we can just ignore the truncate and use the truncate src to do combination.

Alive2 proofs:
https://alive2.llvm.org/ce/z/9D_-qP

[bcl5980]

The motivation of this change comes from https://github.com/rust-lang/rust/issues/107554.

[bcl5980]

rebase with precommit tests.

[spatel]

General proofs with Alive2 are probably not possible for a bitwidth constrained transform like this, but you can include proofs for the specific changes here.
How can we test ctpop?
What happens if the trunc has an extra use?

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4952–4953

Require one less bit for cttz/ctlz when the "is_zero_poison" argument is true?

llvm/test/Transforms/InstCombine/cmp-intrinsic.ll

549

Why truncate to i15?

IIUC, the transform is valid on this example for truncate down to 6 bits, but not 5 bits. So we should have both of those variations.

The source type width should be tested similarly, so if we start with i32, it is ok to truncate to i5, but not i4?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4950

Imo it makes more sense to do this generically.

(icmp P (trunc(X), C)) if KnownBits(X)[OrigWidth:TruncWidth] == 0 just drop the truncate. That will cover all these intrins + any other cases that happen to come up.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4950

I don't want to enable all cases for the trunc. That may cause a lot of potential regressions. By default we should shrink bits if possible not expand it. I prefer to limited the change to the case we can really get improvement.
Maybe I can move this change to SDAG to avoid these concern.

4952–4953

https://alive2.llvm.org/ce/z/pQQavA
It looks one less bit for cttz when is_zero_poison=true is wrong.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4952–4953

ah, I'm wrong. Should be llvm::Log2_32_Ceil(SrcBits - is_zero_poison), not Op0 bitwidth - is_zero_poison

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4950

Seems to only match m_Oneuse(m_Trunc(...)) so seems like only case is the truncate is entirely eliminated, no? Not sure I see how that could cause regressions.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4950

What I mean is not IR level regression. Think about the case source type is i256, the dest type is i32. Remove the trunc and restore the next instruct to be i256 may cause extra instruction on backend.

[bcl5980]

Address comments.

[bcl5980]

add one use negative tests.

[bcl5980]

In D143368#4113564, @spatel wrote:

General proofs with Alive2 are probably not possible for a bitwidth constrained transform like this, but you can include proofs for the specific changes here.
How can we test ctpop?

It looks current code already optimized ctpop. So I just remove the ctpop here.

What happens if the trunc has an extra use?

Line 4920 already limit the code to one-use. And I also add tests for it now.

[bcl5980]

rebase tests.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4950

This is a gray area. Ideally, if we can eliminate IR instructions, we do it, but it if it results in wider instructions, then it's potentially not an improvement for analysis in IR. The codegen concern is secondary, but yes, we do factor that into the decision because sometimes it is not easy to invert the transform.

For this one, it seems unlikely that we'll gain much by using a potentially expensive (in compile-time) ValueTracking API because we probably already have specialized folds for the common patterns where we'd know the high bits are already cleared ('and' or 'lshr').

4954–4955

This needs an explanation comment. Something like:

// If the "is_zero_poison" argument is set, then we know at least 
// one bit is set in the input, so the result is always at least one
// less than the full bitwidth of that input.

4957

Adjust word choice:

// Make sure the destination is wide enough to hold the largest output of the intrinsic.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombinerImpl::foldICmpEquality(ICmpInst &I) {
   if (match(Op1, m_c_And(m_Specific(Op0), m_Value())) &&
       isKnownToBeAPowerOfTwo(Op0, /* OrZero */ false, 0, &I))
     return new ICmpInst(CmpInst::getInversePredicate(Pred), Op1,
                         ConstantInt::getNullValue(Op1->getType()));
 
   return nullptr;
 }
 
-static Instruction *foldICmpWithTrunc(ICmpInst &ICmp,
-                                      InstCombiner::BuilderTy &Builder) {
+Instruction *InstCombinerImpl::foldICmpWithTrunc(ICmpInst &ICmp) {
   ICmpInst::Predicate Pred = ICmp.getPredicate();
   Value *Op0 = ICmp.getOperand(0), *Op1 = ICmp.getOperand(1);
 
   // Try to canonicalize trunc + compare-to-constant into a mask + cmp.
   // The trunc masks high bits while the compare may effectively mask low bits.
   Value *X;
   const APInt *C;
   if (!match(Op0, m_OneUse(m_Trunc(m_Value(X)))) || !match(Op1, m_APInt(C)))
Instruction *InstCombinerImpl::foldICmpWithTrunc(ICmpInst &ICmp) {
   if (Pred == ICmpInst::ICMP_UGT && (~*C).isPowerOf2()) {
     // If C is not-of-power-of-2 (one clear bit):
     // (trunc X) u> C --> (X & (C+1)) == C+1 (are all masked-high-bits set?)
     Constant *MaskC = ConstantInt::get(X->getType(), (*C + 1).zext(SrcBits));
     Value *And = Builder.CreateAnd(X, MaskC);
     return new ICmpInst(ICmpInst::ICMP_EQ, And, MaskC);
   }
 
+  if (auto *II = dyn_cast<IntrinsicInst>(X)) {
+    if (II->getIntrinsicID() == Intrinsic::cttz ||
+        II->getIntrinsicID() == Intrinsic::ctlz) {
+      unsigned MaxRet = SrcBits;
+      if (match(II->getArgOperand(1), m_One()))
+        MaxRet--;
+
+      // Make sure the dest bits is enough to save the intrinsic output's range
+      if (llvm::Log2_32(MaxRet) + 1 <= Op0->getType()->getScalarSizeInBits())
+        if (Instruction *I =
+                foldICmpIntrinsicWithConstant(ICmp, II, C->zext(SrcBits)))
+          return I;
+    }
+  }
+
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::foldICmpWithZextOrSext(ICmpInst &ICmp) {
   assert(isa<CastInst>(ICmp.getOperand(0)) && "Expected cast for operand 0");
   auto *CastOp0 = cast<CastInst>(ICmp.getOperand(0));
   Value *X;
   if (!match(CastOp0, m_ZExtOrSExt(m_Value(X))))
if (CastOp0->getOpcode() == Instruction::PtrToInt &&
     } else if (auto *RHSC = dyn_cast<Constant>(ICmp.getOperand(1))) {
       NewOp1 = ConstantExpr::getIntToPtr(RHSC, SrcTy);
     }
 
     if (NewOp1)
       return new ICmpInst(ICmp.getPredicate(), Op0Src, NewOp1);
   }
 
-  if (Instruction *R = foldICmpWithTrunc(ICmp, Builder))
+  if (Instruction *R = foldICmpWithTrunc(ICmp))
     return R;
 
   return foldICmpWithZextOrSext(ICmp);
 }
 
 static bool isNeutralValue(Instruction::BinaryOps BinaryOp, Value *RHS, bool IsSigned) {
   switch (BinaryOp) {
     default:

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *foldICmpBinOpEqualityWithConstant(ICmpInst &Cmp,
                                                  BinaryOperator *BO,
                                                  const APInt &C);
   Instruction *foldICmpIntrinsicWithConstant(ICmpInst &ICI, IntrinsicInst *II,
                                              const APInt &C);
   Instruction *foldICmpEqIntrinsicWithConstant(ICmpInst &ICI, IntrinsicInst *II,
                                                const APInt &C);
   Instruction *foldICmpBitCast(ICmpInst &Cmp);
+  Instruction *foldICmpWithTrunc(ICmpInst &Cmp);
 
   // Helpers of visitSelectInst().
   Instruction *foldSelectOfBools(SelectInst &SI);
   Instruction *foldSelectExtConst(SelectInst &Sel);
   Instruction *foldSelectOpOp(SelectInst &SI, Instruction *TI, Instruction *FI);
   Instruction *foldSelectIntoOp(SelectInst &SI, Value *, Value *);
   Instruction *foldSPFofSPF(Instruction *Inner, SelectPatternFlavor SPF1,
                             Value *A, Value *B, Instruction &Outer,

llvm/test/Transforms/InstCombine/cmp-intrinsic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 declare i16 @llvm.bswap.i16(i16)
 declare i32 @llvm.bswap.i32(i32)
 declare <2 x i64> @llvm.bswap.v2i64(<2 x i64>)
+declare i32 @llvm.cttz.i32(i32, i1)
 declare i33 @llvm.cttz.i33(i33, i1)
 declare i32 @llvm.ctlz.i32(i32, i1)
 declare i8 @llvm.ctpop.i8(i8)
 declare i11 @llvm.ctpop.i11(i11)
 declare <2 x i32> @llvm.cttz.v2i32(<2 x i32>, i1)
 declare <2 x i32> @llvm.ctlz.v2i32(<2 x i32>, i1)
 declare <2 x i32> @llvm.ctpop.v2i32(<2 x i32>)
 
;
   %tz = tail call i33 @llvm.cttz.i33(i33 %x, i1 false)
   %trunc = trunc i33 %tz to i15
   %cmp = icmp ugt i15 %trunc, 3
   ret i1 %cmp
 }
 
 define i1 @trunc_cttz_ult_other_i33_i15(i33 %x) {
 ; CHECK-LABEL: @trunc_cttz_ult_other_i33_i15(
-; CHECK-NEXT:    [[TZ:%.*]] = tail call i33 @llvm.cttz.i33(i33 [[X:%.*]], i1 false), !range [[RNG1]]
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i33 [[TZ]] to i15
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i15 [[TRUNC]], 7
+; CHECK-NEXT:    [[TMP1:%.*]] = and i33 [[X:%.*]], 127
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i33 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %tz = tail call i33 @llvm.cttz.i33(i33 %x, i1 false)
   %trunc = trunc i33 %tz to i15
   %cmp = icmp ult i15 %trunc, 7
   ret i1 %cmp
 }
 
+define i1 @trunc_cttz_true_ult_other_i32_i5(i32 %x) {
+; CHECK-LABEL: @trunc_cttz_true_ult_other_i32_i5(
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 127
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[TMP1]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %tz = tail call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %trunc = trunc i32 %tz to i5
+  %cmp = icmp ult i5 %trunc, 7
+  ret i1 %cmp
+}
+
+define i1 @trunc_cttz_false_ult_other_i32_i6(i32 %x) {
+; CHECK-LABEL: @trunc_cttz_false_ult_other_i32_i6(
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 127
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[TMP1]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %tz = tail call i32 @llvm.cttz.i32(i32 %x, i1 false)
+  %trunc = trunc i32 %tz to i6
+  %cmp = icmp ult i6 %trunc, 7
+  ret i1 %cmp
+}
+
+
+define i1 @trunc_cttz_false_ult_other_i32_i5(i32 %x) {
+; CHECK-LABEL: @trunc_cttz_false_ult_other_i32_i5(
+; CHECK-NEXT:    [[TZ:%.*]] = tail call i32 @llvm.cttz.i32(i32 [[X:%.*]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[TZ]] to i5
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i5 [[TRUNC]], 7
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %tz = tail call i32 @llvm.cttz.i32(i32 %x, i1 false)
+  %trunc = trunc i32 %tz to i5
+  %cmp = icmp ult i5 %trunc, 7
+  ret i1 %cmp
+}
+
 define i1 @trunc_ctlz_ugt_zero_i32(i32 %x) {
 ; CHECK-LABEL: @trunc_ctlz_ugt_zero_i32(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[X:%.*]], -1
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %trunc = trunc i32 %lz to i15
   %cmp = icmp ugt i15 %trunc, 0
   ret i1 %cmp
 }
 
 define i1 @trunc_ctlz_ugt_one_i32(i32 %x) {
 ; CHECK-LABEL: @trunc_ctlz_ugt_one_i32(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X:%.*]], 1073741824
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %trunc = trunc i32 %lz to i15
   %cmp = icmp ugt i15 %trunc, 1
   ret i1 %cmp
 }
 
-define i1 @trunc_ctlz_ugt_other_i32(i32 %x) {
-; CHECK-LABEL: @trunc_ctlz_ugt_other_i32(
-; CHECK-NEXT:    [[LZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range [[RNG0]]
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[LZ]] to i15
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i15 [[TRUNC]], 4
+define i1 @trunc_ctlz_ugt_other_i32_i15(i32 %x) {
+; CHECK-LABEL: @trunc_ctlz_ugt_other_i32_i15(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X:%.*]], 134217728
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
   %trunc = trunc i32 %lz to i15
   %cmp = icmp ugt i15 %trunc, 4
   ret i1 %cmp
 }
 
+define i1 @trunc_ctlz_true_ugt_other_i32_i5(i32 %x) {
+; CHECK-LABEL: @trunc_ctlz_true_ugt_other_i32_i5(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X:%.*]], 134217728
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 true)
+  %trunc = trunc i32 %lz to i5
+  %cmp = icmp ugt i5 %trunc, 4
+  ret i1 %cmp
+}
+
+define i1 @trunc_ctlz_false_ugt_other_i32_i5(i32 %x) {
+; CHECK-LABEL: @trunc_ctlz_false_ugt_other_i32_i5(
+; CHECK-NEXT:    [[LZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false), !range [[RNG0]]
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[LZ]] to i5
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i5 [[TRUNC]], 4
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
+  %trunc = trunc i32 %lz to i5
+  %cmp = icmp ugt i5 %trunc, 4
+  ret i1 %cmp
+}
+
+define i1 @trunc_ctlz_false_ugt_other_i32_i6(i32 %x) {
+; CHECK-LABEL: @trunc_ctlz_false_ugt_other_i32_i6(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X:%.*]], 134217728
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
+  %trunc = trunc i32 %lz to i6
+  %cmp = icmp ugt i6 %trunc, 4
+  ret i1 %cmp
+}
+
+
 define i1 @trunc_ctpop_eq_zero_i11(i11 %x) {
 ; CHECK-LABEL: @trunc_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)
   %trunc = trunc i11 %pop to i5
   %cmp = icmp eq i5 %trunc, 0