[InstCombine] Match range check pattern with SExt #118910

MatzeB · 2024-12-06T02:01:39Z

= Background

We optimize range check patterns like the following:

  %n_not_negative = icmp sge i32 %n, 0
  call void @llvm.assume(i1 %n_not_negative)
  %a = icmp sge i32 %x, 0
  %b = icmp slt i32 %x, %n
  %c = and i1 %a, %b

to a single unsigned comparison:

  %n_not_negative = icmp sge i32 %n, 0
  call void @llvm.assume(i1 %n_not_negative)
  %c = icmp ult i32 %x, %n

= Extended Pattern

This adds support for a variant of this pattern where the upper range is compared with a sign extended value:

  %n_not_negative = icmp sge i64 %n, 0
  call void @llvm.assume(i1 %n_not_negative)
  %x_sext = sext i32 %x to i64
  %a = icmp sge i32 %x, 0
  %b = icmp slt i64 %x_sext, %n
  %c = and i1 %a, %b

is now optimized to:

  %n_not_negative = icmp sge i64 %n, 0
  call void @llvm.assume(i1 %n_not_negative)
  %x_sext = sext i32 %x to i64
  %c = icmp ult i64 %x_sext, %n

Alive2: https://alive2.llvm.org/ce/z/XVuz9L

= Background LLVM optimizes range check patterns like the following: ``` %nn = and i32 %n, 2147483647 ; %nn is non-negative %a = icmp sge i32 %x, 0 %b = icmp slt i32 %x, %nn %c = and i1 %a, %b ``` to a single unsigned comparison: ``` %nn = and i32 %n, 2147483647 %c = icmp ult i32 %x, %nn ``` = Extended Pattern This adds support for an extended version of this pattern where the upper range is compared with an `sext` value. Example: ``` %nn = and i64 %n, 2147483647 %x_sext = sext i32 %x to i64 %a = icmp sge i32 %x, 0 %b = icmp slt i64 %x_sext, %nn %c = and i1 %a, %b ``` is now optimized to: ``` %nn = and i64 %n, 2147483647 %x_sext = sext i32 %x to i64 %c = icmp ugt i64 %nn, %x_sext ``` In Alive2: https://alive2.llvm.org/ce/z/Ff7KJ_

llvmbot · 2024-12-06T02:22:21Z

@llvm/pr-subscribers-llvm-transforms

Author: Matthias Braun (MatzeB)

Changes

= Background

We optimize range check patterns like the following:

  %nn = and i32 %n, 2147483647   ; %nn is non-negative
  %a = icmp sge i32 %x, 0
  %b = icmp slt i32 %x, %nn
  %c = and i1 %a, %b

to a single unsigned comparison:

  %nn = and i32 %n, 2147483647
  %c = icmp ult i32 %x, %nn

= Extended Pattern

This adds support for a variant of this pattern where the upper range is compared with a sign extended value:

  %nn = and i64 %n, 2147483647
  %x_sext = sext i32 %x to i64
  %a = icmp sge i32 %x, 0
  %b = icmp slt i64 %x_sext, %nn
  %c = and i1 %a, %b

is now optimized to:

  %nn = and i64 %n, 2147483647
  %x_sext = sext i32 %x to i64
  %c = icmp ult i64 %x_sext, %nn

In Alive2: https://alive2.llvm.org/ce/z/Ff7KJ_

Full diff: https://github.com/llvm/llvm-project/pull/118910.diff

2 Files Affected:

(modified) llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp (+17-4)
(modified) llvm/test/Transforms/InstCombine/range-check.ll (+120)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index b4033fc2a418a1..dddab2efd80fb9 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -695,14 +695,27 @@ Value *InstCombinerImpl::simplifyRangeCheck(ICmpInst *Cmp0, ICmpInst *Cmp1,
                                Cmp1->getPredicate());
 
   Value *Input = Cmp0->getOperand(0);
+  Value *Cmp1Op0 = Cmp1->getOperand(0);
+  Value *Cmp1Op1 = Cmp1->getOperand(1);
   Value *RangeEnd;
-  if (Cmp1->getOperand(0) == Input) {
+  if (Cmp1Op0 == Input) {
     // For the upper range compare we have: icmp x, n
-    RangeEnd = Cmp1->getOperand(1);
-  } else if (Cmp1->getOperand(1) == Input) {
+    RangeEnd = Cmp1Op1;
+  } else if (isa<SExtInst>(Cmp1Op0) &&
+             cast<SExtInst>(Cmp1Op0)->getOperand(0) == Input) {
+    // For the upper range compare we have: icmp (sext x), n
+    Input = Cmp1Op0;
+    RangeEnd = Cmp1Op1;
+  } else if (Cmp1Op1 == Input) {
     // For the upper range compare we have: icmp n, x
-    RangeEnd = Cmp1->getOperand(0);
     Pred1 = ICmpInst::getSwappedPredicate(Pred1);
+    RangeEnd = Cmp1Op0;
+  } else if (isa<SExtInst>(Cmp1Op1) &&
+             cast<SExtInst>(Cmp1Op1)->getOperand(0) == Input) {
+    // For the upper range compare we have: icmp n, (sext x)
+    Pred1 = ICmpInst::getSwappedPredicate(Pred1);
+    Input = Cmp1Op1;
+    RangeEnd = Cmp1Op0;
   } else {
     return nullptr;
   }
diff --git a/llvm/test/Transforms/InstCombine/range-check.ll b/llvm/test/Transforms/InstCombine/range-check.ll
index ebb310fb7c1f8f..161fc2d7fdc1dc 100644
--- a/llvm/test/Transforms/InstCombine/range-check.ll
+++ b/llvm/test/Transforms/InstCombine/range-check.ll
@@ -32,6 +32,21 @@ define i1 @test_and1_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_and1_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_and1_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp sge i32 %x, 0
+  %b = icmp slt i64 %x_sext, %nn
+  %c = and i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_and2(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_and2(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -60,6 +75,21 @@ define i1 @test_and2_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_and2_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_and2_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp uge i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp sgt i32 %x, -1
+  %b = icmp sle i64 %x_sext, %nn
+  %c = and i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_and3(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_and3(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -86,6 +116,21 @@ define i1 @test_and3_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_and3_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_and3_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp sgt i64 %nn, %x_sext
+  %b = icmp sge i32 %x, 0
+  %c = and i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_and4(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_and4(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -112,6 +157,21 @@ define i1 @test_and4_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_and4_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_and4_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp uge i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp sge i64 %nn, %x_sext
+  %b = icmp sge i32 %x, 0
+  %c = and i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_or1(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_or1(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -140,6 +200,21 @@ define i1 @test_or1_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_or1_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_or1_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp slt i32 %x, 0
+  %b = icmp sge i64 %x_sext, %nn
+  %c = or i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_or2(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_or2(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -168,6 +243,21 @@ define i1 @test_or2_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_or2_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_or2_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp sle i32 %x, -1
+  %b = icmp sgt i64 %x_sext, %nn
+  %c = or i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_or3(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_or3(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -194,6 +284,21 @@ define i1 @test_or3_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_or3_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_or3_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp sle i64 %nn, %x_sext
+  %b = icmp slt i32 %x, 0
+  %c = or i1 %a, %b
+  ret i1 %c
+}
+
 define i1 @test_or4(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_or4(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -220,6 +325,21 @@ define i1 @test_or4_logical(i32 %x, i32 %n) {
   ret i1 %c
 }
 
+define i1 @test_or4_sext(i32 %x, i64 %n) {
+; CHECK-LABEL: @test_or4_sext(
+; CHECK-NEXT:    [[NN:%.*]] = and i64 [[N:%.*]], 2147483647
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %nn = and i64 %n, 2147483647
+  %x_sext = sext i32 %x to i64
+  %a = icmp slt i64 %nn, %x_sext
+  %b = icmp slt i32 %x, 0
+  %c = or i1 %a, %b
+  ret i1 %c
+}
+
 ; Negative tests
 
 define i1 @negative1(i32 %x, i32 %n) {

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

dtcxzyw

Looks like it also holds for zext: https://alive2.llvm.org/ce/z/vpYSiD

BTW, can you add a proof for (icmp slt x, 0) | (icmp sgt sext(x), n) --> icmp ugt x, n?

MatzeB · 2024-12-09T18:49:22Z

Looks like it also holds for zext: https://alive2.llvm.org/ce/z/vpYSiD

uh oh, the tests should have been and %n, 9223372036854775807. While the rule holds with zext for and i64 %n, 2147483647 where all upper 32bits are zero, the C++ code checks for non-negative number (which is what we actually want) and then zext won't hold any longer.

Let me change the tests to use llvm.assume instead to make things clearer.

MatzeB · 2024-12-09T18:58:08Z

@dtcxzyw

BTW, can you add a proof for (icmp slt x, 0) | (icmp sgt sext(x), n) --> icmp ugt x, n?

What do you mean by "add a proof"? More alive2 examples?

MatzeB · 2024-12-09T19:14:56Z

I added all the unit test examples to the updated alive2 link

dtcxzyw

LGTM. Thank you!

MatzeB requested review from dtcxzyw and goldsteinn December 6, 2024 02:07

MatzeB marked this pull request as ready for review December 6, 2024 02:21

MatzeB requested a review from nikic as a code owner December 6, 2024 02:21

llvmbot added llvm:instcombine Covers the InstCombine, InstSimplify and AggressiveInstCombine passes llvm:transforms labels Dec 6, 2024

dtcxzyw reviewed Dec 6, 2024

View reviewed changes

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp Outdated Show resolved Hide resolved

This was referenced Dec 6, 2024

Task submission dtcxzyw/llvm-opt-benchmark#1312

Open

pre-commit: PR118910 dtcxzyw/llvm-opt-benchmark#1826

Closed

MatzeB added 3 commits December 6, 2024 13:33

Address review feedback

db77078

adjust

2889778

Update InstCombineAndOrXor.cpp

76e6a24

dtcxzyw mentioned this pull request Dec 7, 2024

Fuzz PR118910 dtcxzyw/llvm-mutation-based-fuzz-service#13

Closed

dtcxzyw reviewed Dec 7, 2024

View reviewed changes

Use assume instead of and for the tests

107b0e9

dtcxzyw approved these changes Dec 10, 2024

View reviewed changes

dtcxzyw changed the title ~~Match range check pattern with SExt~~ [InstCombine] Match range check pattern with SExt Dec 10, 2024

MatzeB merged commit e9c68c6 into llvm:main Dec 10, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[InstCombine] Match range check pattern with SExt #118910

[InstCombine] Match range check pattern with SExt #118910

Uh oh!

MatzeB commented Dec 6, 2024 •

edited

Loading

Uh oh!

llvmbot commented Dec 6, 2024

Uh oh!

Uh oh!

dtcxzyw left a comment

Uh oh!

MatzeB commented Dec 9, 2024 •

edited

Loading

Uh oh!

MatzeB commented Dec 9, 2024

Uh oh!

MatzeB commented Dec 9, 2024

Uh oh!

dtcxzyw left a comment

Uh oh!

Uh oh!

Uh oh!

[InstCombine] Match range check pattern with SExt #118910

[InstCombine] Match range check pattern with SExt #118910

Uh oh!

Conversation

MatzeB commented Dec 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Dec 6, 2024

Uh oh!

Uh oh!

dtcxzyw left a comment

Choose a reason for hiding this comment

Uh oh!

MatzeB commented Dec 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MatzeB commented Dec 9, 2024

Uh oh!

MatzeB commented Dec 9, 2024

Uh oh!

dtcxzyw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

MatzeB commented Dec 6, 2024 •

edited

Loading

MatzeB commented Dec 9, 2024 •

edited

Loading