[InstCombine] Add transforms for `(icmp eq/ne (trunc x),(truncOrZext(y)))`->`(icmp eq/ne x,y)`

[goldstein.w.n]

Summary

Three transforms (all commutative):

`(icmp eq/ne (trunc x),(trunc y))`:
    -> if `sizeof(x) == sizeof(y)`:
        -> `(icmp eq/ne x,y)`
            - https://alive2.llvm.org/ce/z/rAj52G
    -> if `sizeof(x) < sizeof(y)`:
        -> `(icmp eq/ne x,(trunc y)`
            - https://alive2.llvm.org/ce/z/vyju55
`(icmp eq/ne (trunc x),(zext y)`
        -> `(icmp eq/ne x,(zext y)`
            - https://alive2.llvm.org/ce/z/mtKZzm

[grandinj]

I think the third line in the commit message should be

-> (icmp eq/ne x,y)

?

[goldstein.w.n]

In D148593#4276624, @grandinj wrote:

I think the third line in the commit message should be

-> (icmp eq/ne x,y)

?

Correct. Fixed summary. Thanks.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

return TrueIfSigned ? new ICmpInst(ICmpInst::ICMP_SLT, ShOp,
                                        ConstantInt::getNullValue(SrcTy))
                         : new ICmpInst(ICmpInst::ICMP_SGT, ShOp,
                                        ConstantInt::getAllOnesValue(SrcTy));
   }
 
   return nullptr;
 }
 
+/// Fold icmp (trunc X), (trunc Y).
+/// Fold icmp (trunc X), zext(Y).
+static Instruction *foldICmpTruncWithTruncOrExt(ICmpInst &Cmp,
+                                                InstCombinerImpl &IC,
+                                                const SimplifyQuery &Q) {
+  if (!Cmp.isEquality())
+    return nullptr;
+
+  Value *X, *Y;
+  ICmpInst::Predicate Pred;
+  bool YIsZext = false;
+  // Try to match icmp (trunc X), (trunc Y)
+  if (match(&Cmp, m_ICmp(Pred, m_Trunc(m_Value(X)), m_Trunc(m_Value(Y))))) {
+    if (X->getType() != Y->getType()) {
+      if (!(Cmp.getOperand(0)->hasOneUse() && Cmp.getOperand(1)->hasOneUse()))
+        return nullptr;
+
+      if (X->getType()->getScalarSizeInBits() >
+          Y->getType()->getScalarSizeInBits())
+        std::swap(X, Y);
+    }
+
+  }
+  // Try to match icmp (trunc X), (zext Y)
+  else if (match(&Cmp,
+                 m_c_ICmp(Pred, m_Trunc(m_Value(X)), m_ZExt(m_Value(Y))))) {
+
+    if (!(isa<ZExtInst>(Cmp.getOperand(0)) ? Cmp.getOperand(0)
+                                           : Cmp.getOperand(1))
+             ->hasOneUse())
+      return nullptr;
+    YIsZext = true;
+  } else {
+    return nullptr;
+  }
+
+  Type *TruncTy, *BaseTy;
+  TruncTy = Cmp.getOperand(0)->getType();
+  BaseTy = X->getType();
+
+  unsigned TruncBits = TruncTy->getScalarSizeInBits();
+  unsigned BaseBits = BaseTy->getScalarSizeInBits();
+
+  // Check if the trunc is unneeded.
+  KnownBits KnownX =
+      computeKnownBits(X, Q.DL, 0, Q.AC, Q.CxtI, Q.DT);
+  if (KnownX.Zero.countl_one() < BaseBits - TruncBits)
+    return nullptr;
+
+  if (!YIsZext) {
+    // If Y is also a trunc, make sure it is unneeded.
+    KnownBits KnownY = computeKnownBits(Y, Q.DL, 0, Q.AC, Q.CxtI, Q.DT);
+    if (KnownY.Zero.countl_one() < BaseBits - TruncBits)
+      return nullptr;
+  }
+
+  Value *NewY = IC.Builder.CreateZExtOrTrunc(Y, BaseTy);
+  return IC.replaceInstUsesWith(
+      Cmp, IC.Builder.CreateICmp(Pred, X, NewY));
+}
+
 /// Fold icmp (xor X, Y), C.
 Instruction *InstCombinerImpl::foldICmpXorConstant(ICmpInst &Cmp,
                                                    BinaryOperator *Xor,
                                                    const APInt &C) {
   if (Instruction *I = foldICmpXorShiftConst(Cmp, Xor, C))
     return I;
 
   Value *X = Xor->getOperand(0);
if (Instruction *Res = foldICmpWithDominatingICmp(I))
     return Res;
 
   if (Instruction *Res = foldICmpUsingBoolRange(I))
     return Res;
 
   if (Instruction *Res = foldICmpUsingKnownBits(I))
     return Res;
 
+  if (Instruction *Res = foldICmpTruncWithTruncOrExt(I, *this, Q))
+    return Res;
+
   // Test if the ICmpInst instruction is used exclusively by a select as
   // part of a minimum or maximum operation. If so, refrain from doing
   // any other folding. This helps out other analyses which understand
   // non-obfuscated minimum and maximum idioms, such as ScalarEvolution
   // and CodeGen. And in this case, at least one of the comparison
   // operands has at least one user besides the compare (the select),
   // which would often largely negate the benefit of folding anyway.
   //

llvm/test/Transforms/InstCombine/eq-of-parts.ll

;
   %c.1 = icmp eq i8 %x.1, %y.1
   %c.2 = icmp eq i8 %x.2, %y.2
   %c.210 = and i1 %c.2, %c.1
   ret i1 %c.210
 }
 
 define i1 @eq_shift_in_zeros(i32 %x, i32 %y) {
 ; CHECK-LABEL: @eq_shift_in_zeros(
-; CHECK-NEXT:    [[X_321:%.*]] = lshr i32 [[X:%.*]], 8
-; CHECK-NEXT:    [[X_1:%.*]] = trunc i32 [[X_321]] to i8
-; CHECK-NEXT:    [[X_32:%.*]] = lshr i32 [[X]], 16
-; CHECK-NEXT:    [[X_2:%.*]] = trunc i32 [[X_32]] to i24
-; CHECK-NEXT:    [[Y_321:%.*]] = lshr i32 [[Y:%.*]], 8
-; CHECK-NEXT:    [[Y_1:%.*]] = trunc i32 [[Y_321]] to i8
-; CHECK-NEXT:    [[Y_32:%.*]] = lshr i32 [[Y]], 16
-; CHECK-NEXT:    [[Y_2:%.*]] = trunc i32 [[Y_32]] to i24
-; CHECK-NEXT:    [[C_1:%.*]] = icmp eq i8 [[X_1]], [[Y_1]]
-; CHECK-NEXT:    [[C_2:%.*]] = icmp eq i24 [[X_2]], [[Y_2]]
-; CHECK-NEXT:    [[C_210:%.*]] = and i1 [[C_2]], [[C_1]]
+; CHECK-NEXT:    [[C_210_UNSHIFTED:%.*]] = xor i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[C_210:%.*]] = icmp ult i32 [[C_210_UNSHIFTED]], 256
 ; CHECK-NEXT:    ret i1 [[C_210]]
 ;
   %x.321 = lshr i32 %x, 8
   %x.1 = trunc i32 %x.321 to i8
   %x.32 = lshr i32 %x, 16
   %x.2 = trunc i32 %x.32 to i24
   %y.321 = lshr i32 %y, 8
   %y.1 = trunc i32 %y.321 to i8
;
   %c.1 = icmp ne i8 %x.1, %y.1
   %c.2 = icmp ne i8 %x.2, %y.2
   %c.210 = or i1 %c.2, %c.1
   ret i1 %c.210
 }
 
 define i1 @ne_shift_in_zeros(i32 %x, i32 %y) {
 ; CHECK-LABEL: @ne_shift_in_zeros(
-; CHECK-NEXT:    [[X_321:%.*]] = lshr i32 [[X:%.*]], 8
-; CHECK-NEXT:    [[X_1:%.*]] = trunc i32 [[X_321]] to i8
-; CHECK-NEXT:    [[X_32:%.*]] = lshr i32 [[X]], 16
-; CHECK-NEXT:    [[X_2:%.*]] = trunc i32 [[X_32]] to i24
-; CHECK-NEXT:    [[Y_321:%.*]] = lshr i32 [[Y:%.*]], 8
-; CHECK-NEXT:    [[Y_1:%.*]] = trunc i32 [[Y_321]] to i8
-; CHECK-NEXT:    [[Y_32:%.*]] = lshr i32 [[Y]], 16
-; CHECK-NEXT:    [[Y_2:%.*]] = trunc i32 [[Y_32]] to i24
-; CHECK-NEXT:    [[C_1:%.*]] = icmp ne i8 [[X_1]], [[Y_1]]
-; CHECK-NEXT:    [[C_2:%.*]] = icmp ne i24 [[X_2]], [[Y_2]]
-; CHECK-NEXT:    [[C_210:%.*]] = or i1 [[C_2]], [[C_1]]
+; CHECK-NEXT:    [[C_210_UNSHIFTED:%.*]] = xor i32 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[C_210:%.*]] = icmp ugt i32 [[C_210_UNSHIFTED]], 255
 ; CHECK-NEXT:    ret i1 [[C_210]]
 ;
   %x.321 = lshr i32 %x, 8
   %x.1 = trunc i32 %x.321 to i8
   %x.32 = lshr i32 %x, 16
   %x.2 = trunc i32 %x.32 to i24
   %y.321 = lshr i32 %y, 8
   %y.1 = trunc i32 %y.321 to i8

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=instcombine < %s | FileCheck %s
 
 declare void @llvm.assume(i1)
 declare void @use(i16)
 define i1 @icmp_trunc_x_trunc_y(i32 %x, i32 %y) {
 ; CHECK-LABEL: @icmp_trunc_x_trunc_y(
 ; CHECK-NEXT:    [[X_LB_ONLY:%.*]] = icmp ult i32 [[X:%.*]], 65536
 ; CHECK-NEXT:    [[Y_LB_ONLY:%.*]] = icmp ult i32 [[Y:%.*]], 65536
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_LB_ONLY]])
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[Y_LB_ONLY]])
-; CHECK-NEXT:    [[X16:%.*]] = trunc i32 [[X]] to i16
-; CHECK-NEXT:    [[Y16:%.*]] = trunc i32 [[Y]] to i16
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i16 [[X16]], [[Y16]]
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[X]], [[Y]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %x_lb_only = icmp ult i32 %x, 65536
   %y_lb_only = icmp ult i32 %y, 65536
   call void @llvm.assume(i1 %x_lb_only)
   call void @llvm.assume(i1 %y_lb_only)
   %x16 = trunc i32 %x to i16
   %y16 = trunc i32 %y to i16
   %r = icmp eq i16 %x16, %y16
   ret i1 %r
 }
 
 define i1 @icmp_trunc_x_trunc_y_2(i32 %x, i64 %y) {
 ; CHECK-LABEL: @icmp_trunc_x_trunc_y_2(
 ; CHECK-NEXT:    [[X_LB_ONLY:%.*]] = icmp ult i32 [[X:%.*]], 12345
 ; CHECK-NEXT:    [[Y_LB_ONLY:%.*]] = icmp ult i64 [[Y:%.*]], 65536
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_LB_ONLY]])
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[Y_LB_ONLY]])
-; CHECK-NEXT:    [[X16:%.*]] = trunc i32 [[X]] to i16
-; CHECK-NEXT:    [[Y16:%.*]] = trunc i64 [[Y]] to i16
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i16 [[Y16]], [[X16]]
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[Y]] to i32
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[TMP1]], [[X]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %x_lb_only = icmp ult i32 %x, 12345
   %y_lb_only = icmp ult i64 %y, 65536
   call void @llvm.assume(i1 %x_lb_only)
   call void @llvm.assume(i1 %y_lb_only)
   %x16 = trunc i32 %x to i16
   %y16 = trunc i64 %y to i16
;
   %r = icmp eq i16 %x16, %y16
   ret i1 %r
 }
 
 define i1 @icmp_trunc_x_zext_y(i32 %x, i8 %y) {
 ; CHECK-LABEL: @icmp_trunc_x_zext_y(
 ; CHECK-NEXT:    [[X_LB_ONLY:%.*]] = icmp ult i32 [[X:%.*]], 65536
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_LB_ONLY]])
-; CHECK-NEXT:    [[X16:%.*]] = trunc i32 [[X]] to i16
-; CHECK-NEXT:    [[Y16:%.*]] = zext i8 [[Y:%.*]] to i16
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i16 [[X16]], [[Y16]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[Y:%.*]] to i32
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[TMP1]], [[X]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %x_lb_only = icmp ult i32 %x, 65536
   call void @llvm.assume(i1 %x_lb_only)
   %x16 = trunc i32 %x to i16
   %y16 = zext i8 %y to i16
   %r = icmp eq i16 %x16, %y16
   ret i1 %r
 }
 
 define i1 @icmp_trunc_x_zext_y_2(i64 %x, i8 %y) {
 ; CHECK-LABEL: @icmp_trunc_x_zext_y_2(
 ; CHECK-NEXT:    [[X_LB_ONLY:%.*]] = icmp ult i64 [[X:%.*]], 65536
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_LB_ONLY]])
-; CHECK-NEXT:    [[X16:%.*]] = trunc i64 [[X]] to i16
-; CHECK-NEXT:    [[Y16:%.*]] = zext i8 [[Y:%.*]] to i16
-; CHECK-NEXT:    [[R:%.*]] = icmp eq i16 [[Y16]], [[X16]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[Y:%.*]] to i64
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i64 [[TMP1]], [[X]]
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %x_lb_only = icmp ult i64 %x, 65536
   call void @llvm.assume(i1 %x_lb_only)
   %x16 = trunc i64 %x to i16
   %y16 = zext i8 %y to i16
   %r = icmp eq i16 %y16, %x16
   ret i1 %r