[InstCombine] Fold (icmp pred (trunc nuw/nsw X), C) -> (icmp pred X, (zext/sext C))

goldsteinn · goldsteinn · commit 1f86aded426e · 2024-05-02T22:49:30.000-05:00
This is valid as long as the sign of the wrap flag doesn't differ from the sign of the `pred`. Proofs: https://alive2.llvm.org/ce/z/35NsrR NB: The online Alive2 hasn't been updated with `trunc nuw/nsw` support, so the proofs must be reproduced locally.
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -1409,6 +1409,19 @@ Instruction *InstCombinerImpl::foldICmpTruncConstant(ICmpInst &Cmp,
                                                      const APInt &C) {
   ICmpInst::Predicate Pred = Cmp.getPredicate();
   Value *X = Trunc->getOperand(0);
+  Type *SrcTy = X->getType();
+  unsigned DstBits = Trunc->getType()->getScalarSizeInBits(),
+           SrcBits = SrcTy->getScalarSizeInBits();
+
+  // Match (icmp pred (trunc nuw/nsw X), C)
+  // Which we can convert to (icmp pred X, (sext/zext C))
+  if (isDesirableIntType(SrcBits) || shouldChangeType(DstBits, SrcBits)) {
+    if (!Cmp.isSigned() && Trunc->hasNoUnsignedWrap())
+      return new ICmpInst(Pred, X, ConstantInt::get(SrcTy, C.zext(SrcBits)));
+    if (Trunc->hasNoSignedWrap())
+      return new ICmpInst(Pred, X, ConstantInt::get(SrcTy, C.sext(SrcBits)));
+  }
+
   if (C.isOne() && C.getBitWidth() > 1) {
     // icmp slt trunc(signum(V)) 1 --> icmp slt V, 1
     Value *V = nullptr;
@@ -1417,10 +1430,6 @@ Instruction *InstCombinerImpl::foldICmpTruncConstant(ICmpInst &Cmp,
                           ConstantInt::get(V->getType(), 1));
   }
 
-  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;
diff --git a/llvm/test/Transforms/InstCombine/cmp-intrinsic.ll b/llvm/test/Transforms/InstCombine/cmp-intrinsic.ll
@@ -620,7 +620,7 @@ define i1 @trunc_cttz_false_ult_other_i32_i6_extra_use(i32 %x) {
 ; CHECK-NEXT:    [[TZ:%.*]] = tail call range(i32 0, 33) i32 @llvm.cttz.i32(i32 [[X:%.*]], i1 false)
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i32 [[TZ]] to i6
 ; CHECK-NEXT:    call void @use6(i6 [[TRUNC]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i6 [[TRUNC]], 7
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[TZ]], 7
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %tz = tail call i32 @llvm.cttz.i32(i32 %x, i1 false)
@@ -722,7 +722,7 @@ define i1 @trunc_ctlz_false_ugt_other_i32_i6_extra_use(i32 %x) {
 ; CHECK-NEXT:    [[LZ:%.*]] = tail call range(i32 0, 33) i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 false)
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i32 [[LZ]] to i6
 ; CHECK-NEXT:    call void @use6(i6 [[TRUNC]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i6 [[TRUNC]], 4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[X]], 134217728
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %lz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 false)
diff --git a/llvm/test/Transforms/InstCombine/icmp-trunc.ll b/llvm/test/Transforms/InstCombine/icmp-trunc.ll
@@ -557,15 +557,9 @@ define i1 @shl1_trunc_sgt4(i32 %a) {
 }
 
 define i1 @eq_nuw(i32 %x) {
-; DL64-LABEL: @eq_nuw(
-; DL64-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 255
-; DL64-NEXT:    [[R:%.*]] = icmp eq i32 [[TMP1]], 123
-; DL64-NEXT:    ret i1 [[R]]
-;
-; DL8-LABEL: @eq_nuw(
-; DL8-NEXT:    [[T:%.*]] = trunc nuw i32 [[X:%.*]] to i8
-; DL8-NEXT:    [[R:%.*]] = icmp eq i8 [[T]], 123
-; DL8-NEXT:    ret i1 [[R]]
+; CHECK-LABEL: @eq_nuw(
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[X:%.*]], 123
+; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nuw i32 %x to i8
   %r = icmp eq i8 %t, 123
@@ -574,8 +568,7 @@ define i1 @eq_nuw(i32 %x) {
 
 define i1 @ult_nuw(i32 %x) {
 ; CHECK-LABEL: @ult_nuw(
-; CHECK-NEXT:    [[T:%.*]] = trunc nuw i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[T]], 45
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i32 [[X:%.*]], 45
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nuw i32 %x to i8
@@ -586,7 +579,7 @@ define i1 @ult_nuw(i32 %x) {
 define i1 @ule_nuw(i32 %x) {
 ; CHECK-LABEL: @ule_nuw(
 ; CHECK-NEXT:    [[T:%.*]] = trunc nuw i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[T]], 46
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i32 [[X]], 46
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
@@ -598,8 +591,7 @@ define i1 @ule_nuw(i32 %x) {
 
 define i1 @ugt_nuw(i32 %x) {
 ; CHECK-LABEL: @ugt_nuw(
-; CHECK-NEXT:    [[T:%.*]] = trunc nuw i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], 12
+; CHECK-NEXT:    [[R:%.*]] = icmp ugt i32 [[X:%.*]], 12
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nuw i32 %x to i8
@@ -610,7 +602,7 @@ define i1 @ugt_nuw(i32 %x) {
 define i1 @uge_nuw(i32 %x) {
 ; CHECK-LABEL: @uge_nuw(
 ; CHECK-NEXT:    [[T:%.*]] = trunc nuw i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], 98
+; CHECK-NEXT:    [[R:%.*]] = icmp ugt i32 [[X]], 98
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
@@ -645,15 +637,9 @@ define i1 @sgt_nuw_fail(i32 %x) {
 }
 
 define i1 @ne_nsw(i32 %x) {
-; DL64-LABEL: @ne_nsw(
-; DL64-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 255
-; DL64-NEXT:    [[R:%.*]] = icmp ne i32 [[TMP1]], 123
-; DL64-NEXT:    ret i1 [[R]]
-;
-; DL8-LABEL: @ne_nsw(
-; DL8-NEXT:    [[T:%.*]] = trunc nsw i32 [[X:%.*]] to i8
-; DL8-NEXT:    [[R:%.*]] = icmp ne i8 [[T]], 123
-; DL8-NEXT:    ret i1 [[R]]
+; CHECK-LABEL: @ne_nsw(
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i32 [[X:%.*]], 123
+; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nsw i32 %x to i8
   %r = icmp ne i8 %t, 123
@@ -662,8 +648,7 @@ define i1 @ne_nsw(i32 %x) {
 
 define i1 @slt_nsw(i32 %x) {
 ; CHECK-LABEL: @slt_nsw(
-; CHECK-NEXT:    [[T:%.*]] = trunc nsw i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp slt i8 [[T]], 45
+; CHECK-NEXT:    [[R:%.*]] = icmp slt i32 [[X:%.*]], 45
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nsw i32 %x to i8
@@ -674,7 +659,7 @@ define i1 @slt_nsw(i32 %x) {
 define i1 @sle_nsw(i16 %x) {
 ; CHECK-LABEL: @sle_nsw(
 ; CHECK-NEXT:    [[T:%.*]] = trunc nsw i16 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp slt i8 [[T]], 46
+; CHECK-NEXT:    [[R:%.*]] = icmp slt i16 [[X]], 46
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
@@ -686,8 +671,7 @@ define i1 @sle_nsw(i16 %x) {
 
 define i1 @sgt_nsw(i32 %x) {
 ; CHECK-LABEL: @sgt_nsw(
-; CHECK-NEXT:    [[T:%.*]] = trunc nsw i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp sgt i8 [[T]], 12
+; CHECK-NEXT:    [[R:%.*]] = icmp sgt i32 [[X:%.*]], 12
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nsw i32 %x to i8
@@ -696,11 +680,17 @@ define i1 @sgt_nsw(i32 %x) {
 }
 
 define i1 @sge_nsw(i64 %x) {
-; CHECK-LABEL: @sge_nsw(
-; CHECK-NEXT:    [[T:%.*]] = trunc nsw i64 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp sgt i8 [[T]], 98
-; CHECK-NEXT:    call void @use(i8 [[T]])
-; CHECK-NEXT:    ret i1 [[R]]
+; DL64-LABEL: @sge_nsw(
+; DL64-NEXT:    [[T:%.*]] = trunc nsw i64 [[X:%.*]] to i8
+; DL64-NEXT:    [[R:%.*]] = icmp sgt i64 [[X]], 98
+; DL64-NEXT:    call void @use(i8 [[T]])
+; DL64-NEXT:    ret i1 [[R]]
+;
+; DL8-LABEL: @sge_nsw(
+; DL8-NEXT:    [[T:%.*]] = trunc nsw i64 [[X:%.*]] to i8
+; DL8-NEXT:    [[R:%.*]] = icmp sgt i8 [[T]], 98
+; DL8-NEXT:    call void @use(i8 [[T]])
+; DL8-NEXT:    ret i1 [[R]]
 ;
   %t = trunc nsw i64 %x to i8
   %r = icmp sge i8 %t, 99
@@ -726,8 +716,7 @@ define i1 @sge_nsw_i48(i48 %x) {
 
 define <2 x i1> @uge_nsw(<2 x i32> %x) {
 ; CHECK-LABEL: @uge_nsw(
-; CHECK-NEXT:    [[T:%.*]] = trunc nsw <2 x i32> [[X:%.*]] to <2 x i8>
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt <2 x i8> [[T]], <i8 44, i8 44>
+; CHECK-NEXT:    [[R:%.*]] = icmp ugt <2 x i32> [[X:%.*]], <i32 44, i32 44>
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %t = trunc nsw <2 x i32> %x to <2 x i8>
