[InstCombine] Fold ashr-exact into a icmp-ugt.

[nadav]

Summary

This commit optimizes the code sequence icmp-XXX (ashr-exact (X, C_1), C_2). Instcombine already implements this optimization for sgt, and this patch adds support to additional predicates. The transformation is legal for all predicates if the 'exact' flag is set, and to SGE, UGE, SLT, ULT when the exact flag is not present.

This pattern is found in the std::vector bounds checks code of the at() method.

Alive2 proof:
https://alive2.llvm.org/ce/z/JT_WL8

[craig.topper]

Upload with full context please.

[lebedev.ri]

Given exact ashr, this is a valid fold regardless of the predicate.
Non-exact is valid for ult/uge / slt/sge predicates.

[nadav]

@lebedev.ri Good catch. Thank you. I verified that when 'exact' is present the transformation is valid for all predicates.

@craig.topper I added context to the patch. Thank you.

[craig.topper]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2242 ↗

(On Diff #399842)

Can ICMP_SGE and ICMP_UGE get here? We're starting from a shift by constant that should be canonicalized to SGT and UGT right?

2242 ↗

(On Diff #399842)

Oops we're starting with a compare with constant that should be canonicalized to SGT to UGT if it was SGE or UGE.

[craig.topper]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2250 ↗

(On Diff #399842)

Does this transform work for ICMP_UGT?

[craig.topper]

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

2219 ↗

(On Diff #399842)

Top of the test context is missing. Should this test be generated by update_test_checks.py?

[nadav]

@craig.topper thank you for the review. I removed the two predicated that are canonicalized before this method is called (ICMP_SGE and ICMP_UGE). I also used the script update_test_checks.py to update the check lines. I also verified with alive2 that the transformation is correct for ICMP_UGT, but only when the exact flag is present (link: https://alive2.llvm.org/ce/z/qBhxAM).

[craig.topper]

In D117252#3242535, @nadav wrote:

@craig.topper thank you for the review. I removed the two predicated that are canonicalized before this method is called (ICMP_SGE and ICMP_UGE). I also used the script update_test_checks.py to update the check lines. I also verified with alive2 that the transformation is correct for ICMP_UGT, but only when the exact flag is present (link: https://alive2.llvm.org/ce/z/qBhxAM).

Thanks Nadav. My UGT question was directed at the transform immediately below. Where it’s looking for SGT and a constant that is 1 less. That would be the form that UGE would canonicalize to I think. Something like this https://alive2.llvm.org/ce/z/6EXVFs

[nadav]

@craig.topper Ah. I think that the ICMP_UGT transformation is correct, but if you don't mind I prefer to land the two changes separately. I can start a second review for the ICMP_UGT change. Is that okay with you?

[craig.topper]

In D117252#3242665, @nadav wrote:

@craig.topper Ah. I think that the ICMP_UGT transformation is correct, but if you don't mind I prefer to land the two changes separately. I can start a second review for the ICMP_UGT change. Is that okay with you?

That's ok with me.

[spatel]

Logic looks good.
Please pre-commit the new tests with baseline CHECK lines.
It would be nice to have at least one test with a vector type and splat constants, so we have coverage for that possibility. I'd also add at least one test with an extra use of the ashr, so we know that doesn't impede the transform.

lib/Transforms/InstCombine/InstCombineCompares.cpp

2243–2244 ↗

(On Diff #399871)

This comment should be updated to be more general. Author's pref for notation/formatting, but it's something like this:

// When ShAmtC can be shifted losslessly:
// (X s>> ShAmtC) < C --> X < (C << ShAmtC)  (either signed/unsigned)
// (X s_exact>> ShAmtC) [AnyPred] C --> X [AnyPred] (C << ShAmtC)

test/Transforms/InstCombine/icmp-shr.ll

716 ↗

(On Diff #399871)

Here and below: remove stale test comments to avoid confusion.

[nadav]

@spatel Thank you for the review. I added the pre-commit tests in a new diff (https://reviews.llvm.org/D117338). I've also made the changes to the comment as you suggested. Are you willing to approve the diff or do you want to see the comment and do another round of review?

[spatel]

In D117252#3244304, @nadav wrote:

@spatel Thank you for the review. I added the pre-commit tests in a new diff (https://reviews.llvm.org/D117338). I've also made the changes to the comment as you suggested. Are you willing to approve the diff or do you want to see the comment and do another round of review?

I don't have anything other than the minor feedback, so no, I don't need another round of review. LGTM.

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2242 ↗

(On Diff #399871)

This should be

if (IsExact || Pred == CmpInst::ICMP_SLT || Pred == CmpInst::ICMP_ULT) {

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

2219 ↗

(On Diff #399871)

Could you please ensure that there is a test for *every* comparison prediocate? (x2 because with/without exact)

[nadav]

@lebedev.ri Okay. I'll add the extra tests, but I'll do it in a separate diff because I've already landed the pre-commit diff that added the extra tests. I also need to look at the optimization that craig pointed out.

[goncharov]

FYI we see internal test became flaky after this change under MSAN. It was resolved later by https://reviews.llvm.org/D117110 (@spatel). Unfortunately I don't have any additional information.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombinerImpl::foldICmpShrConstant(ICmpInst &Cmp,
   bool IsAShr = Shr->getOpcode() == Instruction::AShr;
   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)
+    if (Pred == CmpInst::ICMP_SLT || Pred == CmpInst::ICMP_ULT || IsExact) {
+      // When ShAmtC can be shifted losslessly:
+      // icmp PRED (ashr exact X, ShAmtC), C --> icmp PRED X, (C << ShAmtC)
+      // icmp slt/ult (ashr X, ShAmtC), C --> icmp slt/ult 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() &&

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

 ;
   %s = ashr i8 %x, 3
   %c = icmp slt i8 %s, 10
   ret i1 %c
 }
 
 define i1 @ashr_00_03_noexact(i8 %x) {
 ; CHECK-LABEL: @ashr_00_03_noexact(
-; CHECK-NEXT:    [[S:%.*]] = ashr i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[C:%.*]] = icmp ult i8 [[S]], 10
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i8 [[X:%.*]], 80
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr i8 %x, 3
   %c = icmp ult i8 %s, 10
   ret i1 %c
 }
 
 define i1 @ashr_00_00_exact(i8 %x) {

 ;
   %s = ashr exact i8 %x, 3
   %c = icmp sle i8 %s, 10
   ret i1 %c
 }
 
 define i1 @ashr_00_04_exact(i8 %x) {
 ; CHECK-LABEL: @ashr_00_04_exact(
-; CHECK-NEXT:    [[S:%.*]] = ashr exact i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[C:%.*]] = icmp ugt i8 [[S]], 10
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i8 [[X:%.*]], 80
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr exact i8 %x, 3
   %c = icmp ugt i8 %s, 10
   ret i1 %c
 }
 
 define i1 @ashr_00_05_exact(i8 %x) {
 ; CHECK-LABEL: @ashr_00_05_exact(
-; CHECK-NEXT:    [[S:%.*]] = ashr exact i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[C:%.*]] = icmp ult i8 [[S]], 11
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i8 [[X:%.*]], 88
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr exact i8 %x, 3
   %c = icmp ule i8 %s, 10
   ret i1 %c
 }
 
 define i1 @ashr_00_00_ashr_extra_use(i8 %x, i8* %ptr) {
 ; CHECK-LABEL: @ashr_00_00_ashr_extra_use(
 ; CHECK-NEXT:    [[S:%.*]] = ashr exact i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[C:%.*]] = icmp ult i8 [[S]], 11
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i8 [[X]], 88
 ; CHECK-NEXT:    store i8 [[S]], i8* [[PTR:%.*]], align 1
 ; CHECK-NEXT:    ret i1 [[C]]
 ;
   %s = ashr exact i8 %x, 3
   %c = icmp ule i8 %s, 10
   store i8 %s, i8* %ptr
   ret i1 %c
 }
 
 define <4 x i1> @ashr_00_00_vec(<4 x i8> %x) {
 ; CHECK-LABEL: @ashr_00_00_vec(
-; CHECK-NEXT:    [[S:%.*]] = ashr exact <4 x i8> [[X:%.*]], <i8 3, i8 3, i8 3, i8 3>
-; CHECK-NEXT:    [[C:%.*]] = icmp ult <4 x i8> [[S]], <i8 11, i8 11, i8 11, i8 11>
+; CHECK-NEXT:    [[C:%.*]] = icmp ult <4 x i8> [[X:%.*]], <i8 88, i8 88, i8 88, i8 88>
 ; CHECK-NEXT:    ret <4 x i1> [[C]]
 ;
   %s = ashr exact <4 x i8> %x, <i8 3,i8 3, i8 3, i8 3>
   %c = icmp ule <4 x i8> %s, <i8 10,i8 10,i8 10,i8 10>
   ret <4 x i1> %c
 }
 
 define i1 @lshrult_01_00_exact(i4 %x) {

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

;
   %r = icmp ult i4 %s, 1 ; 0b0001
   ret i1 %r
 }
 
 ; negative test
 
 define i1 @ashr_ult_2(i4 %x) {
 ; CHECK-LABEL: @ashr_ult_2(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 [[X:%.*]], 1
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[S]], 2
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[X:%.*]], 4
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %s = ashr i4 %x, 1
   %r = icmp ult i4 %s, 2 ; 0b0010
   ret i1 %r
 }
 
 ; negative test
 
 define i1 @ashr_ult_3(i4 %x) {
 ; CHECK-LABEL: @ashr_ult_3(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 [[X:%.*]], 1
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[S]], 3
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[X:%.*]], 6
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %s = ashr i4 %x, 1
   %r = icmp ult i4 %s, 3 ; 0b0011
   ret i1 %r
 }
 
 define i1 @ashr_ult_4(i4 %x) {
;
   ret i1 %r
 }
 
 ; negative test
 ; TODO: This is a sign-bit test, but we don't recognize the pattern.
 
 define i1 @ashr_ult_12(i4 %x) {
 ; CHECK-LABEL: @ashr_ult_12(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 [[X:%.*]], 1
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[S]], -4
+; CHECK-NEXT:    [[R:%.*]] = icmp sgt i4 [[X:%.*]], -1
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %s = ashr i4 %x, 1
   %r = icmp ult i4 %s, 12 ; 0b1100
   ret i1 %r
 }
 
 ; negative test
 
 define i1 @ashr_ult_13(i4 %x) {
 ; CHECK-LABEL: @ashr_ult_13(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 [[X:%.*]], 1
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[S]], -3
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[X:%.*]], -6
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %s = ashr i4 %x, 1
   %r = icmp ult i4 %s, 13 ; 0b1101
   ret i1 %r
 }
 
 ; negative test
 
 define i1 @ashr_ult_14(i4 %x) {
 ; CHECK-LABEL: @ashr_ult_14(
-; CHECK-NEXT:    [[S:%.*]] = ashr i4 [[X:%.*]], 1
-; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[S]], -2
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i4 [[X:%.*]], -4
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %s = ashr i4 %x, 1
   %r = icmp ult i4 %s, 14 ; 0b1110
   ret i1 %r
 }
 
 ; negative test, but different transform possible