[InstCombine] try to canonicalize icmp with trunc op into mask and cmp

The motivating test is based on:
https://llvm.org/PR52260

We have better analysis for X == 0, so try harder to form that.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

@@ -4589,6 +4589,31 @@ Instruction *InstCombinerImpl::foldICmpEquality(ICmpInst &I) {
   return nullptr;
 }
 
+static Instruction *foldICmpWithTrunc(ICmpInst &ICmp,
+                                      InstCombiner::BuilderTy &Builder) {
+  const 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_Power2(C))) {
+    if (Pred == ICmpInst::ICMP_ULT) {
+      // (trunc X) u< Pow2C --> (X & MaskC) == 0
+      unsigned SrcBits = X->getType()->getScalarSizeInBits();
+      unsigned DstBits = Op0->getType()->getScalarSizeInBits();
+      APInt MaskC = APInt::getOneBitSet(SrcBits, DstBits) - C->zext(SrcBits);
+      Value *And = Builder.CreateAnd(X, MaskC);
+      Constant *Zero = ConstantInt::getNullValue(X->getType());
+      return new ICmpInst(ICmpInst::ICMP_EQ, And, Zero);
+    }
+    // TODO: Handle ugt.
+  }
+
+  return nullptr;
+}
+
 static Instruction *foldICmpWithZextOrSext(ICmpInst &ICmp,
                                            InstCombiner::BuilderTy &Builder) {
   assert(isa<CastInst>(ICmp.getOperand(0)) && "Expected cast for operand 0");
@@ -4732,6 +4757,9 @@ Instruction *InstCombinerImpl::foldICmpWithCastOp(ICmpInst &ICmp) {
       return new ICmpInst(ICmp.getPredicate(), Op0Src, NewOp1);
   }
 
+  if (Instruction *R = foldICmpWithTrunc(ICmp, Builder))
+    return R;
+
   return foldICmpWithZextOrSext(ICmp, Builder);
 }
 

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

@@ -5,8 +5,8 @@ declare void @use(i8)
 
 define i1 @ult_2(i32 %x) {
 ; CHECK-LABEL: @ult_2(
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 254
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[T]], 2
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
@@ -16,8 +16,8 @@ define i1 @ult_2(i32 %x) {
 
 define <2 x i1> @ult_16_splat(<2 x i16> %x) {
 ; CHECK-LABEL: @ult_16_splat(
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i16> [[X:%.*]] to <2 x i11>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i16> [[X:%.*]], <i16 2032, i16 2032>
-; CHECK-NEXT:    [[R:%.*]] = icmp ult <2 x i11> [[T]], <i11 16, i11 16>
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i16> [[TMP1]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %t = trunc <2 x i16> %x to <2 x i11>
@@ -25,6 +25,8 @@ define <2 x i1> @ult_16_splat(<2 x i16> %x) {
   ret <2 x i1> %r
 }
 
+; negative test - need power-of-2 constant
+
 define i1 @ult_3(i32 %x) {
 ; CHECK-LABEL: @ult_3(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
@@ -36,6 +38,8 @@ define i1 @ult_3(i32 %x) {
   ret i1 %r
 }
 
+; negative test - no extra use allowed
+
 define i1 @ult_2_use(i32 %x) {
 ; CHECK-LABEL: @ult_2_use(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
@@ -53,12 +57,7 @@ define i1 @ult_2_use(i32 %x) {
 
 define i1 @PR52260(i32 %x) {
 ; CHECK-LABEL: @PR52260(
-; CHECK-NEXT:    [[IDXPROM:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    ret i1 true
-; CHECK-NEXT:    [[IDX:%.*]] = getelementptr inbounds [3 x i32], [3 x i32]* @a, i64 0, i64 [[IDXPROM]]
-; CHECK-NEXT:    [[T1:%.*]] = load i32, i32* [[IDX]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[T1]] to i8
-; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp ult i8 [[TMP1]], 2
-; CHECK-NEXT:    ret i1 [[TOBOOL]]
 ;
   %idxprom = sext i32 %x to i64
   %idx = getelementptr inbounds [3 x i32], [3 x i32]* @a, i64 0, i64 %idxprom
@@ -69,4 +68,3 @@ define i1 @PR52260(i32 %x) {
   %tobool = icmp eq i8 %conv2, 0
   ret i1 %tobool
 }
-

llvm/test/Transforms/InstCombine/signed-truncation-check.ll

@@ -397,8 +397,8 @@ define i1 @positive_trunc_signbit_logical(i32 %arg) {
 
 define i1 @positive_trunc_base(i32 %arg) {
 ; CHECK-LABEL: @positive_trunc_base(
-; CHECK-NEXT:    [[T1:%.*]] = trunc i32 [[ARG:%.*]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[ARG:%.*]], 65408
-; CHECK-NEXT:    [[T5_SIMPLIFIED:%.*]] = icmp ult i16 [[T1]], 128
+; CHECK-NEXT:    [[T5_SIMPLIFIED:%.*]] = icmp eq i32 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[T5_SIMPLIFIED]]
 ;
   %t1 = trunc i32 %arg to i16
@@ -411,8 +411,8 @@ define i1 @positive_trunc_base(i32 %arg) {
 
 define i1 @positive_trunc_base_logical(i32 %arg) {
 ; CHECK-LABEL: @positive_trunc_base_logical(
-; CHECK-NEXT:    [[T1:%.*]] = trunc i32 [[ARG:%.*]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[ARG:%.*]], 65408
-; CHECK-NEXT:    [[T5_SIMPLIFIED:%.*]] = icmp ult i16 [[T1]], 128
+; CHECK-NEXT:    [[T5_SIMPLIFIED:%.*]] = icmp eq i32 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[T5_SIMPLIFIED]]
 ;
   %t1 = trunc i32 %arg to i16