[InstCombine] improve folds for icmp gt/lt (shr X, C1), C2

[spatel]

Summary

We can always eliminate the shift in: icmp gt/lt (shr X, C1), C2 --> icmp gt/lt X, C'
This patch was supposed to just be an efficiency improvement because we were doing this 3-step process to fold:

IC: Visiting:   %c = icmp ugt i4 %s, 1
IC: ADD:   %s = lshr i4 %x, 1
IC: ADD:   %1 = udiv i4 %x, 2
IC: Old =   %c = icmp ugt i4 %1, 1
    New =   <badref> = icmp uge i4 %x, 4
IC: ADD:   %c = icmp uge i4 %x, 4
IC: ERASE   %2 = icmp ugt i4 %1, 1
IC: Visiting:   %c = icmp uge i4 %x, 4
IC: Old =   %c = icmp uge i4 %x, 4
    New =   <badref> = icmp ugt i4 %x, 3
IC: ADD:   %c = icmp ugt i4 %x, 3
IC: ERASE   %2 = icmp uge i4 %x, 4
IC: Visiting:   %c = icmp ugt i4 %x, 3
IC: DCE:   %1 = udiv i4 %x, 2
IC: ERASE   %1 = udiv i4 %x, 2
IC: DCE:   %s = lshr i4 %x, 1
IC: ERASE   %s = lshr i4 %x, 1
IC: Visiting:   ret i1 %c

When we could go directly to canonical icmp form:

IC: Visiting:   %c = icmp ugt i4 %s, 1
IC: Old =   %c = icmp ugt i4 %s, 1
    New =   <badref> = icmp ugt i4 %x, 3
IC: ADD:   %c = icmp ugt i4 %x, 3
IC: ERASE   %1 = icmp ugt i4 %s, 1
IC: ADD:   %s = lshr i4 %x, 1
IC: DCE:   %s = lshr i4 %x, 1
IC: ERASE   %s = lshr i4 %x, 1
IC: Visiting:   %c = icmp ugt i4 %x, 3

...but then I noticed that the folds were incomplete too:
https://godbolt.org/g/aB2hLE

Here are attempts to prove the logic with Alive:
https://rise4fun.com/Alive/92o

Name: lshr_ult
Pre: ((C2 << C1) u>> C1) == C2
%sh = lshr i8 %x, C1
%r = icmp ult i8 %sh, C2
  =>
%r = icmp ult i8 %x, (C2 << C1)

Name: ashr_slt
Pre: ((C2 << C1) >> C1) == C2
%sh = ashr i8 %x, C1
%r = icmp slt i8 %sh, C2
  =>
%r = icmp slt i8 %x, (C2 << C1)

Name: lshr_ugt
Pre: (C2 != 255) && (((C2+1) << C1) u>> C1) == (C2+1)
%sh = lshr i8 %x, C1
%r = icmp ugt i8 %sh, C2
  =>
%r = icmp ugt i8 %x, ((C2+1) << C1) - 1

Name: ashr_sgt
Pre: (C2 != 127) && ((C2+1) << C1 != -128) && (((C2+1) << C1) >> C1) == (C2+1)
%sh = ashr i8 %x, C1
%r = icmp sgt i8 %sh, C2
  =>
%r = icmp sgt i8 %x, ((C2+1) << C1) - 1

We did something similar for 'shl' in D28406.

[craig.topper]

I think this changes behavior of the current code with "exact" shifts. Did we not have tests for that?

For example:

define i1 @lshrugt_01_01_exact(i4 %x) {
; CHECK-LABEL: @lshrugt_01_01(
; CHECK-NEXT:    [[C:%.*]] = icmp ugt i4 %x, 2
; CHECK-NEXT:    ret i1 [[C]]
;
  %s = lshr exact i4 %x, 1
  %c = icmp ugt i4 %s, 1
  ret i1 %c
}

define i1 @lshrugt_01_01(i4 %x) {
; CHECK-LABEL: @lshrugt_01_01(
; CHECK-NEXT:    [[C:%.*]] = icmp ugt i4 %x, 3
; CHECK-NEXT:    ret i1 [[C]]
;
  %s = lshr i4 %x, 1
  %c = icmp ugt i4 %s, 1
  ret i1 %c
}

[spatel]

In D38514#887862, @craig.topper wrote:

I think this changes behavior of the current code with "exact" shifts. Did we not have tests for that?

Nice catch - let me add some handling for that. And no, there are no tests that expose the difference (there's only one candidate 'exact.ll' test for this fold currently, but the constant didn't wiggle). Before rL314837, you could have completely bogus logic for the fold (because I did of course) and not break anything.

[spatel]

Patch updated:

1. Account for exact shifts in the same way as the existing code. It's not clear to me that using exact universally produces a better result, but I'm not proposing to change that here. In rL314907, I duplicated the entire range of tests with exactness to prove that nothing is changing for any of those tests.

2. Rearrange the code to try to make the relationship between all of these variants clearer. It's only the non-exact + greater-than cases where we have to inc/dec the constant. All other cases are just a plain shl of the compare constant.

Here's Alive code for the exact cases:

Name: ashr_exact_sgt
Pre: ((C2 << C1) >> C1) == C2
%sh = ashr exact i8 %x, C1
%r = icmp sgt i8 %sh, C2
  =>
%r = icmp sgt i8 %x, (C2 << C1)

Name: ashr_exact_slt
Pre: ((C2 << C1) >> C1) == C2
%sh = ashr exact i8 %x, C1
%r = icmp slt i8 %sh, C2
  =>
%r = icmp slt i8 %x, (C2 << C1)

Name: lshr_exact_ugt
Pre: ((C2 << C1) u>> C1) == C2
%sh = lshr exact i8 %x, C1
%r = icmp ugt i8 %sh, C2
  =>
%r = icmp ugt i8 %x, (C2 << C1)

Name: lshr_exact_ult
Pre: ((C2 << C1) u>> C1) == C2
%sh = lshr exact i8 %x, C1
%r = icmp ult i8 %sh, C2
  =>
%r = icmp ult i8 %x, (C2 << C1)

https://rise4fun.com/Alive/Erp

[craig.topper]

LGTM

[spatel]

Given the possibility raised in PR34838 ( D38591 ), I'm going to change the asserts to checks here.

llvm/trunk/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombiner::foldICmpShrConstant(ICmpInst &Cmp,
   // Check that the shift amount is in range. If not, don't perform undefined
   // shifts. When the shift is visited it will be simplified.
   unsigned TypeBits = C.getBitWidth();
   unsigned ShAmtVal = ShiftAmt->getLimitedValue(TypeBits);
   if (ShAmtVal >= TypeBits || ShAmtVal == 0)
     return nullptr;
 
   bool IsAShr = Shr->getOpcode() == Instruction::AShr;
-  if (!Cmp.isEquality()) {
-    // If we have an unsigned comparison and an ashr, we can't simplify this.
-    // Similarly for signed comparisons with lshr.
-    if (Cmp.isSigned() != IsAShr)
-      return nullptr;
+  bool IsExact = Shr->isExact();
+  Type *ShrTy = Shr->getType();
+  // TODO: If we could guarantee that InstSimplify would handle all of the
+  // constant-value-based preconditions in the folds below, then we could assert
+  // those conditions rather than checking them. This is difficult because of
+  // undef/poison (PR34838).
+  if (IsAShr) {
+    if (Pred == CmpInst::ICMP_SLT || (Pred == CmpInst::ICMP_SGT && IsExact)) {
+      // icmp slt (ashr X, ShAmtC), C --> icmp slt X, (C << ShAmtC)
+      // icmp sgt (ashr exact X, ShAmtC), C --> icmp sgt X, (C << ShAmtC)
+      APInt ShiftedC = C.shl(ShAmtVal);
+      if (ShiftedC.ashr(ShAmtVal) == C)
+        return new ICmpInst(Pred, X, ConstantInt::get(ShrTy, ShiftedC));
+    }
+    if (Pred == CmpInst::ICMP_SGT) {
+      // icmp sgt (ashr X, ShAmtC), C --> icmp sgt X, ((C + 1) << ShAmtC) - 1
+      APInt ShiftedC = (C + 1).shl(ShAmtVal) - 1;
+      if (!C.isMaxSignedValue() && !(C + 1).shl(ShAmtVal).isMinSignedValue() &&
+          (ShiftedC + 1).ashr(ShAmtVal) == (C + 1))
+        return new ICmpInst(Pred, X, ConstantInt::get(ShrTy, ShiftedC));
+    }
+  } else {
+    if (Pred == CmpInst::ICMP_ULT || (Pred == CmpInst::ICMP_UGT && IsExact)) {
+      // icmp ult (lshr X, ShAmtC), C --> icmp ult X, (C << ShAmtC)
+      // icmp ugt (lshr exact X, ShAmtC), C --> icmp ugt X, (C << ShAmtC)
+      APInt ShiftedC = C.shl(ShAmtVal);
+      if (ShiftedC.lshr(ShAmtVal) == C)
+        return new ICmpInst(Pred, X, ConstantInt::get(ShrTy, ShiftedC));
+    }
+    if (Pred == CmpInst::ICMP_UGT) {
+      // icmp ugt (lshr X, ShAmtC), C --> icmp ugt X, ((C + 1) << ShAmtC) - 1
+      APInt ShiftedC = (C + 1).shl(ShAmtVal) - 1;
+      if ((ShiftedC + 1).lshr(ShAmtVal) == (C + 1))
+        return new ICmpInst(Pred, X, ConstantInt::get(ShrTy, ShiftedC));
+    }
+  }
 
-    // Otherwise, all lshr and most exact ashr's are equivalent to a udiv/sdiv
-    // by a power of 2.  Since we already have logic to simplify these,
-    // transform to div and then simplify the resultant comparison.
-    if (IsAShr && (!Shr->isExact() || ShAmtVal == TypeBits - 1))
-      return nullptr;
+  if (!Cmp.isEquality())
 
-    // Revisit the shift (to delete it).
-    Worklist.Add(Shr);
-
-    Constant *DivCst = ConstantInt::get(
-        Shr->getType(), APInt::getOneBitSet(TypeBits, ShAmtVal));
-
-    Value *Tmp = IsAShr ? Builder.CreateSDiv(X, DivCst, "", Shr->isExact())
-                        : Builder.CreateUDiv(X, DivCst, "", Shr->isExact());
-
-    Cmp.setOperand(0, Tmp);
-
-    // If the builder folded the binop, just return it.
-    BinaryOperator *TheDiv = dyn_cast<BinaryOperator>(Tmp);
-    if (!TheDiv)
-      return &Cmp;
-
-    // Otherwise, fold this div/compare.
-    assert(TheDiv->getOpcode() == Instruction::SDiv ||
-           TheDiv->getOpcode() == Instruction::UDiv);
-
-    Instruction *Res = foldICmpDivConstant(Cmp, TheDiv, C);
-    assert(Res && "This div/cst should have folded!");
-    return Res;
-  }
-
   // Handle equality comparisons of shift-by-constant.
 
   // If the comparison constant changes with the shift, the comparison cannot
   // succeed (bits of the comparison constant cannot match the shifted value).
   // This should be known by InstSimplify and already be folded to true/false.
   assert(((IsAShr && C.shl(ShAmtVal).ashr(ShAmtVal) == C) ||
           (!IsAShr && C.shl(ShAmtVal).lshr(ShAmtVal) == C)) &&
          "Expected icmp+shr simplify did not occur.");
 
   // Check if the bits shifted out are known to be zero. If so, we can compare
   // against the unshifted value:
   //  (X & 4) >> 1 == 2  --> (X & 4) == 4.
-  Constant *ShiftedCmpRHS = ConstantInt::get(Shr->getType(), C << ShAmtVal);
+  Constant *ShiftedCmpRHS = ConstantInt::get(ShrTy, C << ShAmtVal);
   if (Shr->hasOneUse()) {
     if (Shr->isExact())
       return new ICmpInst(Pred, X, ShiftedCmpRHS);
 
     // Otherwise strength reduce the shift into an 'and'.
     APInt Val(APInt::getHighBitsSet(TypeBits, TypeBits - ShAmtVal));
-    Constant *Mask = ConstantInt::get(Shr->getType(), Val);
+    Constant *Mask = ConstantInt::get(ShrTy, Val);
     Value *And = Builder.CreateAnd(X, Mask, Shr->getName() + ".mask");
     return new ICmpInst(Pred, And, ShiftedCmpRHS);
   }
 
   return nullptr;
 }
 
 /// Fold icmp (udiv X, Y), C.

llvm/trunk/test/Transforms/InstCombine/icmp-shr-lt-gt.ll

 ;
   %s = lshr i4 %x, 3
   %c = icmp ult i4 %s, 15
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_00(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_00(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], 0
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, 1
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 0
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_01(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_01(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], 1
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, 3
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 1
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_02(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_02(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, 5
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 2
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_03(i4 %x) {

 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 11
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_12(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_12(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], -4
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, -7
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 12
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_13(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_13(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], -3
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, -5
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 13
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_14(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_14(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], -2
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, -3
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 14
   ret i1 %c
 }
 
 define i1 @ashrsgt_01_15(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_01_15(
 ; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, -1
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp sgt i4 %s, 15
   ret i1 %c
 }
 
 define i1 @ashrsgt_02_00(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_02_00(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 2
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], 0
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, 3
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 2
   %c = icmp sgt i4 %s, 0
   ret i1 %c
 }
 
 define i1 @ashrsgt_02_01(i4 %x) {

 ;
   %s = ashr i4 %x, 2
   %c = icmp sgt i4 %s, 13
   ret i1 %c
 }
 
 define i1 @ashrsgt_02_14(i4 %x) {
 ; CHECK-LABEL: @ashrsgt_02_14(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 2
-; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 [[S]], -2
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i4 %x, -5
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 2
   %c = icmp sgt i4 %s, 14
   ret i1 %c
 }
 
 define i1 @ashrsgt_02_15(i4 %x) {

 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 0
   ret i1 %c
 }
 
 define i1 @ashrslt_01_01(i4 %x) {
 ; CHECK-LABEL: @ashrslt_01_01(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], 1
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, 2
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 1
   ret i1 %c
 }
 
 define i1 @ashrslt_01_02(i4 %x) {
 ; CHECK-LABEL: @ashrslt_01_02(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, 4
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 2
   ret i1 %c
 }
 
 define i1 @ashrslt_01_03(i4 %x) {
 ; CHECK-LABEL: @ashrslt_01_03(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], 3
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, 6
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 3
   ret i1 %c
 }
 
 define i1 @ashrslt_01_04(i4 %x) {

 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 12
   ret i1 %c
 }
 
 define i1 @ashrslt_01_13(i4 %x) {
 ; CHECK-LABEL: @ashrslt_01_13(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], -3
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, -6
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 13
   ret i1 %c
 }
 
 define i1 @ashrslt_01_14(i4 %x) {
 ; CHECK-LABEL: @ashrslt_01_14(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], -2
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, -4
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 14
   ret i1 %c
 }
 
 define i1 @ashrslt_01_15(i4 %x) {
 ; CHECK-LABEL: @ashrslt_01_15(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 1
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], -1
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, -2
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 1
   %c = icmp slt i4 %s, 15
   ret i1 %c
 }
 
 define i1 @ashrslt_02_00(i4 %x) {
 ; CHECK-LABEL: @ashrslt_02_00(
 ; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, 0
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 2
   %c = icmp slt i4 %s, 0
   ret i1 %c
 }
 
 define i1 @ashrslt_02_01(i4 %x) {
 ; CHECK-LABEL: @ashrslt_02_01(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 2
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], 1
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, 4
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 2
   %c = icmp slt i4 %s, 1
   ret i1 %c
 }
 
 define i1 @ashrslt_02_02(i4 %x) {

 ;
   %s = ashr i4 %x, 2
   %c = icmp slt i4 %s, 14
   ret i1 %c
 }
 
 define i1 @ashrslt_02_15(i4 %x) {
 ; CHECK-LABEL: @ashrslt_02_15(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 %x, 2
-; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 [[S]], -1
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i4 %x, -4
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i4 %x, 2
   %c = icmp slt i4 %s, 15
   ret i1 %c
 }
 
 define i1 @ashrslt_03_00(i4 %x) {