[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

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

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

Does this transform work for ICMP_UGT?

[craig.topper]

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

2219

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

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

This should be

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

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

2219

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.

lib/Transforms/InstCombine/InstCombineCompares.cpp

   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)) {
+    if (Pred == CmpInst::ICMP_SGE || Pred == CmpInst::ICMP_UGE ||
+        Pred == CmpInst::ICMP_SLT || Pred == CmpInst::ICMP_ULT || 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));
     }
 
     // If the compare constant has significant bits above the lowest sign-bit,

   // instruction that we're folding.
   bool ProdOV = (DivIsSigned ? Prod.sdiv(*C2) : Prod.udiv(*C2)) != C;
 
   ICmpInst::Predicate Pred = Cmp.getPredicate();
 
   // If the division is known to be exact, then there is no remainder from the
   // divide, so the covered range size is unit, otherwise it is the divisor.
   APInt RangeSize = Div->isExact() ? APInt(C2->getBitWidth(), 1) : *C2;
 

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

 ; CHECK-LABEL: @lshrugt_03_15_exact(
 ; CHECK-NEXT:    ret i1 false
 ;
   %s = lshr exact i4 %x, 3
   %c = icmp ugt i4 %s, 15
   ret i1 %c
 }
 
+define i1 @lashr_00_00_noexact(i8 %x) {
+; CHECK-LABEL: @lashr_00_00_noexact(
+; CHECK-NEXT:    %c = icmp sgt i8 %x, 79
+;
+  %s = ashr i8 %x, 3
+  %c = icmp sge i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_01_noexact(i8 %x) {
+; CHECK-LABEL: @lashr_00_01_noexact(
+; CHECK-NEXT:    %c = icmp slt i8 %x, 80
+;
+  %s = ashr i8 %x, 3
+  %c = icmp slt i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_03_noexact(i8 %x) {
+; CHECK-LABEL: @lashr_00_03_noexact(
+; CHECK-NEXT:    %c = icmp ult i8 %x, 80
+;
+  %s = ashr i8 %x, 3
+  %c = icmp ult i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_00_exact(i8 %x) {
+; CHECK-LABEL: @lashr_00_00_exact(
+; CHECK-NEXT:    %c = icmp ne i8 %x, 80
+;
+  %s = ashr exact i8 %x, 3
+  %c = icmp ne i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_01_exact(i8 %x) {
+; CHECK-LABEL: @lashr_00_01_exact(
+; CHECK-NEXT:    %c = icmp eq i8 %x, 80
+;
+  %s = ashr exact i8 %x, 3
+  %c = icmp eq i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_02_exact(i8 %x) {
+; CHECK-LABEL: @lashr_00_02_exact(
+; CHECK-NEXT:    %c = icmp sgt i8 %x, 80
+;
+  %s = ashr exact i8 %x, 3
+  %c = icmp sgt i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_03_exact(i8 %x) {
+; CHECK-LABEL: @lashr_00_03_exact(
+; CHECK-NEXT:    %c = icmp slt i8 %x, 88
+;
+  %s = ashr exact i8 %x, 3
+  %c = icmp sle i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_04_exact(i8 %x) {
+; CHECK-LABEL: @lashr_00_04_exact(
+; CHECK-NEXT:    %c = icmp ugt i8 %x, 80
+;
+  %s = ashr exact i8 %x, 3
+  %c = icmp ugt i8 %s, 10
+  ret i1 %c
+}
+
+define i1 @lashr_00_05_exact(i8 %x) {
+; CHECK-LABEL: @lashr_00_05_exact(
+; CHECK-NEXT:    %c = icmp ult i8 %x, 88
+;
+  %s = ashr exact i8 %x, 3
+  %c = icmp ule i8 %s, 10
+  ret i1 %c
+}
+
 define i1 @lshrult_01_00_exact(i4 %x) {
 ; CHECK-LABEL: @lshrult_01_00_exact(
 ; CHECK-NEXT:    ret i1 false
 ;
   %s = lshr exact i4 %x, 1
   %c = icmp ult i4 %s, 0
   ret i1 %c
 }