[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

   // 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_SLT || (Pred == CmpInst::ICMP_SGT && IsExact) ||
+        (Pred == CmpInst::ICMP_UGT && IsExact)) {
       // icmp slt (ashr X, ShAmtC), C --> icmp slt X, (C << ShAmtC)
       // icmp sgt (ashr exact X, ShAmtC), C --> icmp sgt X, (C << ShAmtC)
+      // icmp ugt (ashr exact X, ShAmtC), C --> icmp ugt X, (C << ShAmtC)
       APInt ShiftedC = C.shl(ShAmtVal);
       if (ShiftedC.ashr(ShAmtVal) == C)
         return new ICmpInst(Pred, X, ConstantInt::get(ShrTy, ShiftedC));

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

   ret i1 %c
 }
 
+define i1 @lashrugt_00_00_exact(i8 %x) {
+; CHECK-LABEL: @lashrugt_00_00_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 @lshrult_01_00_exact(i4 %x) {
 ; CHECK-LABEL: @lshrult_01_00_exact(
 ; CHECK-NEXT:    ret i1 false