[InstCombine] Fold nuw left-shifts in `ugt`/`ule` comparisons.

[bryant]

Summary

This transforms

%a = shl nuw %x, c1
%b = icmp {ugt|ule} %a, c0

into

%b = icmp {ugt|ule} %x, (c0 >> c1)

z3:

(declare-const x (_ BitVec 64))
(declare-const c0 (_ BitVec 64))
(declare-const c1 (_ BitVec 64))

(push)
(assert (= x (bvlshr (bvshl x c1) c1)))  ; nuw
(assert (not (= (bvugt (bvshl x c1) c0)
                (bvugt x
                       (bvlshr c0 c1)))))
(check-sat)
(get-model)
(pop)

(push)
(assert (= x (bvlshr (bvshl x c1) c1)))  ; nuw
(assert (not (= (bvule (bvshl x c1) c0)
                (bvule x
                       (bvlshr c0 c1)))))
(check-sat)
(get-model)
(pop)

[spatel]

The patch will crash for vectors:
define <2 x i1> @icmp_ugt_16_vec(<2 x i64> %x) {

%c = shl nuw <2 x i64> %x, <i64 16, i64 16>
%d = icmp ugt <2 x i64> %c, <i64 65535, i64 65535>
ret <2 x i1> %d

}

Splat vectors should be handled the same as scalars. You can look at the fold just above yours for a model of how to fix it.

lib/Transforms/InstCombine/InstCombineCompares.cpp

1990–1991

Nit: start sentence with capital letter.

test/Transforms/InstCombine/icmp-shl-zext-simplify.ll

3–13 ↗

(On Diff #75607)

1. You can simplify the tests in all cases by:

a. Remove the leading zext
b. Add 'nuw' to the shl
c. Remove the trailing zext (just return the i1)

2. Please use the python script at utils/update_test_checks.py to auto-generate the check lines.

3. I'd prefer to see tests that check different constant values rather than different type sizes. Ie, we can be pretty sure that InstCombine folds work for arbitrary sizes, so those tests don't add a lot of value, but we need to make sure that the new constant is created correctly for the new fold:

define i1 @icmp_ugt_16trunc(i64 %x) {
  %shl = shl nuw i64 %x, 16
  %cmp = icmp ugt i64 %shl, 65537
  ret i1 %cmp
}

[bryant]

• Handle the vector case properly.
• Remove redundant tests.
• Test vector and non-zero cases.

[bryant]

test/Transforms/InstCombine/icmp-shl-zext-simplify.ll

37 ↗

(On Diff #75631)

This trunc bothers me. On X86, I think it generates an extra pand:

.LCPI3_0:
        .short  65535                   # 0xffff
        .short  65535                   # 0xffff
        .short  65535                   # 0xffff
        .short  0                       # 0x0
        .short  65535                   # 0xffff
        .short  65535                   # 0xffff
        .short  65535                   # 0xffff
        .short  0                       # 0x0
        .text
        .globl  icmp_ule_i64x2
        .p2align        4, 0x90
        .type   icmp_ule_i64x2,@function
icmp_ule_i64x2:                         # @icmp_ule_i64x2
        .cfi_startproc
# BB#0:
        pand    .LCPI3_0(%rip), %xmm0      # <=======
        pxor    %xmm1, %xmm1
        pcmpeqd %xmm0, %xmm1
        pshufd  $177, %xmm1, %xmm0      # xmm0 = xmm1[1,0,3,2]
        pand    %xmm1, %xmm0
        retq

The same goes for icmp_ule_64, but no extra X86 op is generated in that case.

[bryant]

• Handle uge, ult.
• Prioritize over case that outputs trunc.

[spatel]

test/Transforms/InstCombine/icmp-shl-zext-simplify.ll

37 ↗

(On Diff #75631)

I noticed that example too and agree it doesn't look like a good fold. We shouldn't be creating illegal types in InstCombine without a fallback in a later pass to make sure the backend is not harmed.

But this is the trunc fold above yours that is firing, right? Ie, this test does not change with your patch.
Please add a vector test that is affected by your patch.

You can add all of the tests with the current output based on trunk, then rebase this patch after that commit, so we'll see just the diffs in the resulting IR.

[bryant]

test/Transforms/InstCombine/icmp-shl-zext-simplify.ll

37 ↗

(On Diff #75631)

I've already moved this patch above the trunc fold, though.

[bryant]

Add tests to handle uge (icmp_uge_8x2) and ult (icmp_ult_8). Also add vector test directly affected by this patch (icmp_ugt_16x2).

[spatel]

test/Transforms/InstCombine/icmp-shl-zext-simplify.ll

37 ↗

(On Diff #75631)

Your last update came in while I was writing my last comment, so I hadn't seen it.
Now there's even more reason to add the tests with current trunk output before the code patch because we should document the current behavior that your patch is going to alter. Please do that and rebase.

[bryant]

Rebase onto D25955.

[spatel]

In D25913#578829, @bryant wrote:

Rebase onto D25955.

Thanks for the preliminary patches.

As you can see from the current trunk transforms, we always canonicalize the compare predicates to ugt/ult. Unless you see a case where that doesn't happen, I think you can adjust this patch to only fire when we have ugt/ult rather than ugt/ule and remove the "pre-transform" step.

[bryant]

Saved a few slocs.

[bryant]

Merge uge/ule with ugt/ult, respectively.

[spatel]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1962–1963

Why return a ULE icmp here rather than ULT?

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1962–1963

Are you suggesting to emit x u< [(c - 1) >> s] - 1? That just seems messier.

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1962–1963

+ 1, not - 1.

[spatel]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1955

">u" should be ">=u"

1961

I don't think you need to check that the constant is ugt(0) here because:
icmp ult %x, 0
must always be simplified to 'false' ahead of this.
Let me know if I've misunderstood that check.

1962–1963

If you don't change the predicate, I think it makes the code easier to read (as well as short-circuiting the subsequent canonicalization to ULT).

Could look more like this?

if (Shl->hasNoUnsignedWrap() &&
    (Pred == ICmpInst::ICMP_UGT || Pred == ICmpInst::ICMP_ULT)) {
  Type *CTy = IntegerType::get(Cmp.getContext(), C->getBitWidth());
  if (X->getType()->isVectorTy())
    CTy = VectorType::get(CTy, X->getType()->getVectorNumElements());

  APInt ShiftedC = Pred == ICmpInst::ICMP_ULT ? (*C - 1).lshr(*ShiftAmt) + 1
                                              : C->lshr(*ShiftAmt);
  return new ICmpInst(Pred, X, ConstantInt::get(CTy, ShiftedC));
}

test/Transforms/InstCombine/icmp-shl-zext-simplify.ll

39–42 ↗

(On Diff #75793)

Please change the constant in this test so we have more test coverage to show the correct non-zero constant for a ule/ult comparison. Eg, 196608 (0x30000) --> 4 ?

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1955

So, just to recap: This transformation only works for all c when ugt and ule (cf. the z3 in my opening post) and fails for uge:

z3
(push)
(declare-const s (_ BitVec 8))
(declare-const c (_ BitVec 8))
(declare-const x (_ BitVec 8))

(assert (= (bvlshr (bvshl x s) s) x))

(assert (not (= (bvuge (bvshl x s) c) (bvuge x (bvlshr c s)))))
(check-sat)
(get-model)
(pop)

sat
(model 
  (define-fun s () (_ BitVec 8)
    #x10)
  (define-fun x () (_ BitVec 8)
    #x00)
  (define-fun c () (_ BitVec 8)
    #x01)
)

But the original rule can be transformed from ugt to uge:

(x << s) ugt c = x ugt (c >> s)
(x << s) uge (c + 1) = x ugt (c >> s) if c < -1
(x << s) uge (c + 1) = x uge (c >> s) + 1 if c < -1
(x << s) uge c = x uge ((c - 1) >> s) + 1 if c > 0

The exact same holds for ule to ult, but c has to be non-zero.

Maybe I'll add this derivation to the comments.

1961

It's a safety check that doesn't seem so expensive. If C is zero by some chance for whatever reason, then this transformation is incorrect and really mustn't happen. Are you really sure that I should remove it? Better safe than slow.

[spatel]

Yes, adding the derivation to the comments would be good.

I wasn't really concerned about the speed of the extra check for *C != 0; it just simplified the code to not have the extra check for that half of the transform. Since I've gotten that sort of assumption wrong before, I'm cc'ing David and Eli to see if they have a better idea.

[bryant]

• replaced non-zero check for c with note in comments.
• added derivation of ult case in comments.
• use icmp-shl-nuw.ll instead of icmp-shl-zext-simplify.ll
• added non-zero constant vector and non-pow2 shift test cases.

[efriedma]

"X <u 0" should be reliably folded to "false" before you reach this point, if that's the question... InstSimplify always runs before any other InstCombine transform. Maybe add an assert to be on the safe side.

Not sure about "X >u 0" -> "X != 0"; I think that actually happens later, although I guess it doesn't matter in this context.

lib/Transforms/InstCombine/InstCombineCompares.cpp

1963

"// (X << S) <u (C + 1) = X <u (C >> S) + 1 if C < - 1"

This comment is confusing... when exactly is "C < - 1" true?

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1963

-1 == all ones. Assuming that addition wraps, X <=u C is equivalent to X <=u (C + 1) if C + 1 doesn't wrap, or if C is less than all ones.

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1963

Replace the second <=u with <u.

[efriedma]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1963

Ah, okay. You could probably write that in a way that's a little more obvious.

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1963

C < -1u? C <u -1? C <u ~0u?

[efriedma]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1963

I would probably go with "C <u -1u", but any of those is fine.

[bryant]

• Comment fixes. -1 => ~0u.
• Add sanity assertion for ult.

[bryant]

Sanjay: Would you mind landing this if there are no further objections? Thanks in advance. (:

[spatel]

In D25913#584760, @bryant wrote:

Sanjay: Would you mind landing this if there are no further objections? Thanks in advance. (:

Sure - let's ping Eli to make sure there are no other suggestions, but the code looks good to me.

test/Transforms/InstCombine/icmp-shl-nuw.ll

54–55

"16x2" --> "12x2" ?

[efriedma]

Just one minor comment.

lib/Transforms/InstCombine/InstCombineCompares.cpp

1960

Isn't CTy always equal to "X->getType()"?

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1960

Since we're morphing (x << s) ugt c into x ugt (c >> s), CTy needs to be enough to hold c >> s.

c's type would be sufficient thanks to the right shift.

But unlike other binops, shift nodes place no constraint between the types of shiftee and shift amount (I'm really just inferring from my memory of TargetSelectionDAG.td, so I could be wrong).

So X-getType() might not be wide enough to hold c >> s.

1970

ShiftC. Ugh.

test/Transforms/InstCombine/icmp-shl-nuw.ll

54–55

Thanks for catching this.

[efriedma]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1960

In IR the two operands of an shl always have the same type. (The rules change once you get to SelectionDAG.)

[bryant]

• Re-worded the main comment to flow better.
• ShiftC => ShiftedC.
• Fix test case name.

[bryant]

Simplify since ICmp implies identically typed lhs + rhs.

[bryant]

lib/Transforms/InstCombine/InstCombineCompares.cpp

1960

You're right, just checked LangRef. And on top of that, icmp obviously needs both sides to be the same. Fixed.

[spatel]

In D25913#584987, @bryant wrote:

Simplify since ICmp implies identically typed lhs + rhs.

Yes - that looks nicer. Unlike the DAG, IR has actual documentation that we can reference. :)
"Both arguments to the ‘shl‘ instruction must be the same integer or vector of integer type."
http://llvm.org/docs/LangRef.html#shl-instruction

Sorry for the earlier bogus suggestion to grab the code from the trunc fold.

lib/Transforms/InstCombine/InstCombineCompares.cpp

if (Shl->hasOneUse() && isSignBitCheck(Pred, *C, TrueIfSigned)) {
     Constant *Mask = ConstantInt::get(
         X->getType(),
         APInt::getOneBitSet(TypeBits, TypeBits - ShiftAmt->getZExtValue() - 1));
     Value *And = Builder->CreateAnd(X, Mask, Shl->getName() + ".mask");
     return new ICmpInst(TrueIfSigned ? ICmpInst::ICMP_NE : ICmpInst::ICMP_EQ,
                         And, Constant::getNullValue(And->getType()));
   }
 
+  // When the shift is nuw and pred is >u (for all C) or <u (for non-zero C),
+  // comparison only really happens in the pre-shifted bits. This also holds for
+  // <=u and >u, but the latter two are canonicalized into the former.
+  if (Shl->hasNoUnsignedWrap() &&
+      (Pred == ICmpInst::ICMP_UGT || Pred == ICmpInst::ICMP_ULT)) {
+    Type *CTy = IntegerType::get(Cmp.getContext(), C->getBitWidth());
+    if (X->getType()->isVectorTy())
+      CTy = VectorType::get(CTy, X->getType()->getVectorNumElements());
+    // Derivation for the ult case:
+    // (X << S) <=u C is equiv to X <=u (C >> S) for all C
+    // (X << S) <u (C + 1) is equiv to X <u (C >> S) + 1 if C <u ~0u
+    // (X << S) <u C is equiv to X <u ((C - 1) >> S) + 1 if C >u 0
+    assert((Pred != ICmpInst::ICMP_ULT || C->ugt(0)) &&
+           "Encountered `ult 0` that should have been eliminated by "
+           "InstSimplify.");
+    APInt ShiftedC = Pred == ICmpInst::ICMP_ULT ? (*C - 1).lshr(*ShiftAmt) + 1
+                                                : C->lshr(*ShiftAmt);
+    return new ICmpInst(Pred, X, ConstantInt::get(CTy, ShiftC));
+  }
+
   // Transform (icmp pred iM (shl iM %v, N), C)
   // -> (icmp pred i(M-N) (trunc %v iM to i(M-N)), (trunc (C>>N))
   // Transform the shl to a trunc if (trunc (C>>N)) has no loss and M-N.
   // This enables us to get rid of the shift in favor of a trunc which can be
   // free on the target. It has the additional benefit of comparing to a
   // smaller constant, which will be target friendly.
   unsigned Amt = ShiftAmt->getLimitedValue(TypeBits - 1);
   if (Shl->hasOneUse() && Amt != 0 && C->countTrailingZeros() >= Amt &&
       DL.isLegalInteger(TypeBits - Amt)) {
     Type *TruncTy = IntegerType::get(Cmp.getContext(), TypeBits - Amt);
     if (X->getType()->isVectorTy())
       TruncTy = VectorType::get(TruncTy, X->getType()->getVectorNumElements());
     Constant *NewC =
         ConstantInt::get(TruncTy, C->ashr(*ShiftAmt).trunc(TypeBits - Amt));
     return new ICmpInst(Pred, Builder->CreateTrunc(X, TruncTy), NewC);
   }
 
   return nullptr;
 }
 
 /// Fold icmp ({al}shr X, Y), C.
 Instruction *InstCombiner::foldICmpShrConstant(ICmpInst &Cmp,
                                                BinaryOperator *Shr,
                                                const APInt *C) {
   // An exact shr only shifts out zero bits, so:
   // icmp eq/ne (shr X, Y), 0 --> icmp eq/ne X, 0
   Value *X = Shr->getOperand(0);

test/Transforms/InstCombine/icmp-shl-nuw.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -instcombine -S | FileCheck %s
 
 define i1 @icmp_ugt_32(i64) {
 ; CHECK-LABEL: @icmp_ugt_32(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw i64 %0, 32
-; CHECK-NEXT:    [[D:%.*]] = icmp ugt i64 [[C]], 4294967295
+; CHECK-NEXT:    [[D:%.*]] = icmp ne i64 %0, 0
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %c = shl nuw i64 %0, 32
   %d = icmp ugt i64 %c, 4294967295
   ret i1 %d
 }
 
 define i1 @icmp_ule_64(i128) {
 ; CHECK-LABEL: @icmp_ule_64(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw i128 %0, 64
-; CHECK-NEXT:    [[D:%.*]] = icmp ult i128 [[C]], 18446744073709551616
+; CHECK-NEXT:    [[D:%.*]] = icmp eq i128 %0, 0
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %c = shl nuw i128 %0, 64
   %d = icmp ule i128 %c, 18446744073709551615
   ret i1 %d
 }
 
 define i1 @icmp_ugt_16(i64) {
 ; CHECK-LABEL: @icmp_ugt_16(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw i64 %0, 16
-; CHECK-NEXT:    [[D:%.*]] = icmp ugt i64 [[C]], 1048575
+; CHECK-NEXT:    [[D:%.*]] = icmp ugt i64 %0, 15
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %c = shl nuw i64 %0, 16
   %d = icmp ugt i64 %c, 1048575 ; 0x0f_ffff
   ret i1 %d
 }
 
 define <2 x i1> @icmp_ule_16x2(<2 x i64>) {
 ; CHECK-LABEL: @icmp_ule_16x2(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw <2 x i64> %0, <i64 16, i64 16>
-; CHECK-NEXT:    [[D:%.*]] = icmp ult <2 x i64> [[C]], <i64 65536, i64 65536>
+; CHECK-NEXT:    [[D:%.*]] = icmp eq <2 x i64> %0, zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[D]]
 ;
   %c = shl nuw <2 x i64> %0, <i64 16, i64 16>
   %d = icmp ule <2 x i64> %c, <i64 65535, i64 65535>
   ret <2 x i1> %d
 }
 
+define <2 x i1> @icmp_ule_16x2_nonzero(<2 x i64>) {
+; CHECK-LABEL: @icmp_ule_16x2_nonzero(
+; CHECK-NEXT:    [[D:%.*]] = icmp ult <2 x i64> %0, <i64 4, i64 4>
+; CHECK-NEXT:    ret <2 x i1> [[D]]
+;
+  %c = shl nuw <2 x i64> %0, <i64 16, i64 16>
+  %d = icmp ule <2 x i64> %c, <i64 196608, i64 196608>  ; 0x03_0000
+  ret <2 x i1> %d
+}
+
+define <2 x i1> @icmp_ule_16x2_non_pow2_shift(<2 x i64>) {
+; CHECK-LABEL: @icmp_ule_16x2_non_pow2_shift(
+; CHECK-NEXT:    [[D:%.*]] = icmp ult <2 x i64> %0, <i64 4, i64 4>
+; CHECK-NEXT:    ret <2 x i1> [[D]]
+;
+  %c = shl nuw <2 x i64> %0, <i64 12, i64 12>
+  %d = icmp ule <2 x i64> %c, <i64 12288, i64 12288>  ; 0x3000
+  ret <2 x i1> %d
+}
+
 define i1 @icmp_ult_8(i64) {
 ; CHECK-LABEL: @icmp_ult_8(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw i64 %0, 8
-; CHECK-NEXT:    [[D:%.*]] = icmp ult i64 [[C]], 4095
+; CHECK-NEXT:    [[D:%.*]] = icmp ult i64 %0, 16
 ; CHECK-NEXT:    ret i1 [[D]]
 ;
   %c = shl nuw i64 %0, 8
   %d = icmp ult i64 %c, 4095 ; 0x0fff
   ret i1 %d
 }
 
 define <2 x i1> @icmp_uge_8x2(<2 x i16>) {
 ; CHECK-LABEL: @icmp_uge_8x2(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw <2 x i16> %0, <i16 8, i16 8>
-; CHECK-NEXT:    [[D:%.*]] = icmp ugt <2 x i16> [[C]], <i16 4094, i16 4094>
+; CHECK-NEXT:    [[D:%.*]] = icmp ugt <2 x i16> %0, <i16 15, i16 15>
 ; CHECK-NEXT:    ret <2 x i1> [[D]]
 ;
   %c = shl nuw <2 x i16> %0, <i16 8, i16 8>
   %d = icmp uge <2 x i16> %c, <i16 4095, i16 4095>
   ret <2 x i1> %d
 }
 
 define <2 x i1> @icmp_ugt_16x2(<2 x i32>) {
 ; CHECK-LABEL: @icmp_ugt_16x2(
-; CHECK-NEXT:    [[C:%.*]] = shl nuw <2 x i32> %0, <i32 16, i32 16>
-; CHECK-NEXT:    [[D:%.*]] = icmp ugt <2 x i32> [[C]], <i32 1048575, i32 1048575>
+; CHECK-NEXT:    [[D:%.*]] = icmp ugt <2 x i32> %0, <i32 15, i32 15>
 ; CHECK-NEXT:    ret <2 x i1> [[D]]
 ;
   %c = shl nuw <2 x i32> %0, <i32 16, i32 16>
   %d = icmp ugt <2 x i32> %c, <i32 1048575, i32 1048575>
   ret <2 x i1> %d
 }