[InstCombine] canonicalize icmp with trunc op into mask and cmp, part 2

[spatel]

Summary

If C is a high-bit mask:
(trunc X) u< C --> (X & C) != C (are any masked-high-bits clear?)

This extends the fold added with:
acabad9ff6bf (https://alive2.llvm.org/ce/z/aFr7qV)

We discussed using decomposeBitTestICmp() to generalize this, but that function doesn't line up with the other fold that I was imagining (maybe there's some way to adapt/invert the logic?).

This patch also modifies the code to create the mask constant from the earlier patch in an attempt to make the bit-masking relationships clearer. I can make that an NFC pre-commit to be safer.

Here are Alive2 generalizations the folds:
https://alive2.llvm.org/ce/z/u-ZpC_ (the previous patch)
https://alive2.llvm.org/ce/z/YsuAu2 (ult this patch)

https://alive2.llvm.org/ce/z/ekktQP (ugt low bitmask)
https://alive2.llvm.org/ce/z/pJY9wR (ugt one clear bit)

[spatel]

Patch updated:
I am including the sibling pair of 'ugt' folds here, so we can try to see if there's a cleaner way to implement this.
decomposeBitTestICmp() assumes the new compare constant is zero, so we could use that as-is for 2 of these, but then we'd still have the other 2 folds here.
I'm not sure if generalizing that helper would be useful for any other potential callers.

[spatel]

These are the Alive2 proofs for the 'ugt' folds:
https://alive2.llvm.org/ce/z/ekktQP
https://alive2.llvm.org/ce/z/pJY9wR

[spatel]

Patch updated:
Fixed a code comment that was copy-pasted without being updated to match the actual logic in the code.

[lebedev.ri]

Like i asked in mail, this really seems like something for decomposeBitTestICmp().

[spatel]

In D112634#3106893, @lebedev.ri wrote:

Like i asked in mail, this really seems like something for decomposeBitTestICmp().

I don't disagree, but we have 2 at least problems:

1. It doesn't cover the pair of folds that result in a non-zero constant for the new icmp.
2. It catches signbit tests that may currently get folded in the opposite direction (see inline comment), so we'll need to remove a transform.

So I'm leaning towards adding these as the first step. Reversing the existing fold could lead to regressions in IR or codegen, so that seems more risky, and I'll potentially have to chase down regressions to get that to stick.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1897–1899

This is an opposing transform if we use decomposeBitMask to make this patch more general.
This was added with:
dfa3b0954145ec6

[spatel]

I noticed that there are a pair of similar folds for icmp (add X, C1), C just ahead of the fold in D113366. I think we're missing something similar to what I'm proposing here, but I haven't worked out the bitmask relationships.

[spatel]

Ping.
As I mentioned, I'd prefer not to try to add folds and invert an existing fold in one patch. I did add more instcombine tests for a follow-up patch that would use decomposeBitTestICmp.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4642

I'm guessing we prefer inequality check over relational? https://alive2.llvm.org/ce/z/9mn4Tx

4646

Are these supposed to be in this patch? I'm not seeing alive2 links in the description for them.

[lebedev.ri]

LGTM.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

4642

Yes, we should canonicalize to equality on that pattern:
https://alive2.llvm.org/ce/z/MWK88Q

4646

Yes, I put the 'ugt' patterns in for symmetry (but I can commit in pieces if that seems better). The alive2 links were in a subsequent comment, but I just added them to the description.

[fhahn]

It looks like this introduced a regression with -Oz: https://github.com/llvm/llvm-project/issues/53321

It would be great if you could take a look.

[MatzeB]

It seems this leads to slightly worse x86 codegen. Where we previously had:

%v1 = trunc i64 %v0 to i32
%cmp = icmp slt i32 %v1, 0

resulting in:

testl   %eax, %eax
js      .LBB0_1

instcombine now changes this to:

%v1 = and i64 %v0, 2147483648
%cmp = icmp eq i64 %v1, 0

resulting in

testl   $-2147483648, %eax              # imm = 0x80000000
jne     .LBB0_2

Though I guess we best fix this by adding more x86 patterns...

[spatel]

In D112634#3392128, @fhahn wrote:

It looks like this introduced a regression with -Oz: https://github.com/llvm/llvm-project/issues/53321

It would be great if you could take a look.

Sorry I missed this when it was sent - the code looks fine now, so I closed the bug.

[spatel]

In D112634#3965504, @MatzeB wrote:

It seems this leads to slightly worse x86 codegen. Where we previously had:

%v1 = trunc i64 %v0 to i32
%cmp = icmp slt i32 %v1, 0

resulting in:

testl   %eax, %eax
js      .LBB0_1

instcombine now changes this to:

%v1 = and i64 %v0, 2147483648
%cmp = icmp eq i64 %v1, 0

resulting in

testl   $-2147483648, %eax              # imm = 0x80000000
jne     .LBB0_2

Though I guess we best fix this by adding more x86 patterns...

Yes, we can fix this up at some point in codegen/isel. For any mask+cmp of an i64/i32/i16 try to convert to a signbit test of a smaller power-of-2 type to avoid using a constant in the test instruction?

[spatel]

In D112634#3966800, @spatel wrote:

Yes, we can fix this up at some point in codegen/isel. For any mask+cmp of an i64/i32/i16 try to convert to a signbit test of a smaller power-of-2 type to avoid using a constant in the test instruction?

Proposal to do this target-independently (but using target hooks to limit it):
https://reviews.llvm.org/D139363

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombinerImpl::foldICmpAndConstant(ICmpInst &Cmp,
   // X & -C != -C -> X <= u ~C
   //   iff C is a power of 2
   if (Cmp.getOperand(1) == Y && C.isNegatedPowerOf2()) {
     auto NewPred =
         Pred == CmpInst::ICMP_EQ ? CmpInst::ICMP_UGT : CmpInst::ICMP_ULE;
     return new ICmpInst(NewPred, X, SubOne(cast<Constant>(Cmp.getOperand(1))));
   }
 
   // (X & C2) == 0 -> (trunc X) >= 0
   // (X & C2) != 0 -> (trunc X) <  0
   //   iff C2 is a power of 2 and it masks the sign bit of a legal integer type.
   const APInt *C2;
   if (And->hasOneUse() && C.isZero() && match(Y, m_APInt(C2))) {
     int32_t ExactLogBase2 = C2->exactLogBase2();
     if (ExactLogBase2 != -1 && DL.isLegalInteger(ExactLogBase2 + 1)) {
       Type *NTy = IntegerType::get(Cmp.getContext(), ExactLogBase2 + 1);
       if (auto *AndVTy = dyn_cast<VectorType>(And->getType()))
         NTy = VectorType::get(NTy, AndVTy->getElementCount());
       Value *Trunc = Builder.CreateTrunc(X, NTy);
static Instruction *foldICmpWithTrunc(ICmpInst &ICmp,
   Value *X;
   const APInt *C;
   if (!match(Op0, m_OneUse(m_Trunc(m_Value(X)))) || !match(Op1, m_APInt(C)))
     return nullptr;
 
   unsigned SrcBits = X->getType()->getScalarSizeInBits();
   if (Pred == ICmpInst::ICMP_ULT) {
     if (C->isPowerOf2()) {
-      // If C is a power-of-2:
+      // If C is a power-of-2 (one set bit):
       // (trunc X) u< C --> (X & -C) == 0 (are all masked-high-bits clear?)
       Constant *MaskC = ConstantInt::get(X->getType(), (-*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 C is negative-power-of-2.
+    // If C is a negative power-of-2 (high-bit mask):
+    // (trunc X) u< C --> (X & C) != C (are any masked-high-bits clear?)
+    if (C->isNegatedPowerOf2()) {
+      Constant *MaskC = ConstantInt::get(X->getType(), C->zext(SrcBits));
+      Value *And = Builder.CreateAnd(X, MaskC);
+      return new ICmpInst(ICmpInst::ICMP_NE, And, MaskC);
+    }
+  }
+
+  if (Pred == ICmpInst::ICMP_UGT) {
+    // If C is a low-bit-mask (C+1 is a power-of-2):
+    // (trunc X) u> C --> (X & ~C) != 0 (are any masked-high-bits set?)
+    if (C->isMask()) {
+      Constant *MaskC = ConstantInt::get(X->getType(), (~*C).zext(SrcBits));
+      Value *And = Builder.CreateAnd(X, MaskC);
+      Constant *Zero = ConstantInt::getNullValue(X->getType());
+      return new ICmpInst(ICmpInst::ICMP_NE, And, Zero);
+    }
+    // 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?)
+    if ((~*C).isPowerOf2()) {
+      Constant *MaskC = ConstantInt::get(X->getType(), (*C + 1).zext(SrcBits));
+      Value *And = Builder.CreateAnd(X, MaskC);
+      return new ICmpInst(ICmpInst::ICMP_EQ, And, MaskC);
+    }
   }
-  // TODO: Handle ugt.
 
   return nullptr;
 }
 
 static Instruction *foldICmpWithZextOrSext(ICmpInst &ICmp,
                                            InstCombiner::BuilderTy &Builder) {
   assert(isa<CastInst>(ICmp.getOperand(0)) && "Expected cast for operand 0");
   auto *CastOp0 = cast<CastInst>(ICmp.getOperand(0));

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

;
   %t2 = trunc i32 %conv1 to i8
   %conv2 = and i8 %t2, 127
   %tobool = icmp eq i8 %conv2, 0
   ret i1 %tobool
 }
 
 define i1 @ult_192(i32 %x) {
 ; CHECK-LABEL: @ult_192(
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[T]], -64
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 192
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i32 [[TMP1]], 192
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   %r = icmp ult i8 %t, 192 ; 0b1100_0000
   ret i1 %r
 }
 
 define <2 x i1> @ult_2044_splat(<2 x i16> %x) {
 ; CHECK-LABEL: @ult_2044_splat(
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i16> [[X:%.*]] to <2 x i11>
-; CHECK-NEXT:    [[R:%.*]] = icmp ult <2 x i11> [[T]], <i11 -4, i11 -4>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i16> [[X:%.*]], <i16 2044, i16 2044>
+; CHECK-NEXT:    [[R:%.*]] = icmp ne <2 x i16> [[TMP1]], <i16 2044, i16 2044>
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %t = trunc <2 x i16> %x to <2 x i11>
   %r = icmp ult <2 x i11> %t, <i11 2044, i11 2044> ; 0b111_1111_1100
   ret <2 x i1> %r
 }
 
+; negative test - need high-bit-mask constant
+
 define i1 @ult_96(i32 %x) {
 ; CHECK-LABEL: @ult_96(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
 ; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[T]], 96
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   %r = icmp ult i8 %t, 96  ; 0b0110_0000
   ret i1 %r
 }
 
+; negative test - no extra use allowed
+
 define i1 @ult_192_use(i32 %x) {
 ; CHECK-LABEL: @ult_192_use(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    [[R:%.*]] = icmp ult i8 [[T]], -64
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   call void @use(i8 %t)
   %r = icmp ult i8 %t, 192
   ret i1 %r
 }
 
 define i1 @ugt_3(i32 %x) {
 ; CHECK-LABEL: @ugt_3(
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], 3
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 252
+; CHECK-NEXT:    [[R:%.*]] = icmp ne i32 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   %r = icmp ugt i8 %t, 3
   ret i1 %r
 }
 
 define <2 x i1> @ugt_7_splat(<2 x i16> %x) {
 ; CHECK-LABEL: @ugt_7_splat(
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i16> [[X:%.*]] to <2 x i11>
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt <2 x i11> [[T]], <i11 7, i11 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i16> [[X:%.*]], <i16 2040, i16 2040>
+; CHECK-NEXT:    [[R:%.*]] = icmp ne <2 x i16> [[TMP1]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %t = trunc <2 x i16> %x to <2 x i11>
   %r = icmp ugt <2 x i11> %t, <i11 7, i11 7>
   ret <2 x i1> %r
 }
 
+; negative test - need low-bit-mask constant
+
 define i1 @ugt_4(i32 %x) {
 ; CHECK-LABEL: @ugt_4(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
 ; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], 4
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   %r = icmp ugt i8 %t, 4
   ret i1 %r
 }
 
+; negative test - no extra use allowed
+
 define i1 @ugt_3_use(i32 %x) {
 ; CHECK-LABEL: @ugt_3_use(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], 3
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   call void @use(i8 %t)
   %r = icmp ugt i8 %t, 3
   ret i1 %r
 }
 
 define i1 @ugt_253(i32 %x) {
 ; CHECK-LABEL: @ugt_253(
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], -3
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[X:%.*]], 254
+; CHECK-NEXT:    [[R:%.*]] = icmp eq i32 [[TMP1]], 254
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   %r = icmp ugt i8 %t, 253
   ret i1 %r
 }
 
 define <2 x i1> @ugt_2043_splat(<2 x i16> %x) {
 ; CHECK-LABEL: @ugt_2043_splat(
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i16> [[X:%.*]] to <2 x i11>
-; CHECK-NEXT:    [[R:%.*]] = icmp ugt <2 x i11> [[T]], <i11 -5, i11 -5>
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i16> [[X:%.*]], <i16 2044, i16 2044>
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i16> [[TMP1]], <i16 2044, i16 2044>
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %t = trunc <2 x i16> %x to <2 x i11>
   %r = icmp ugt <2 x i11> %t, <i11 2043, i11 2043> ; 0b111_1111_101
   ret <2 x i1> %r
 }
 
+; negative test - need not-of-power-of-2 constant
+
 define i1 @ugt_252(i32 %x) {
 ; CHECK-LABEL: @ugt_252(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
 ; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], -4
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8
   %r = icmp ugt i8 %t, 252
   ret i1 %r
 }
 
+; negative test - no extra use allowed
+
 define i1 @ugt_253_use(i32 %x) {
 ; CHECK-LABEL: @ugt_253_use(
 ; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[X:%.*]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[T]])
 ; CHECK-NEXT:    [[R:%.*]] = icmp ugt i8 [[T]], -3
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %t = trunc i32 %x to i8