Skip to content

Conversation

MatzeB
Copy link
Contributor

@MatzeB MatzeB commented Dec 6, 2024

= 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_
@MatzeB MatzeB requested review from dtcxzyw and goldsteinn December 6, 2024 02:07
@MatzeB MatzeB marked this pull request as ready for review December 6, 2024 02:21
@MatzeB MatzeB requested a review from nikic as a code owner December 6, 2024 02:21
@llvmbot llvmbot added llvm:instcombine Covers the InstCombine, InstSimplify and AggressiveInstCombine passes llvm:transforms labels Dec 6, 2024
@llvmbot
Copy link
Member

llvmbot commented Dec 6, 2024

@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) {

Copy link
Member

@dtcxzyw dtcxzyw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Contributor Author

MatzeB commented Dec 9, 2024

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
Copy link
Contributor Author

MatzeB commented Dec 9, 2024

@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
Copy link
Contributor Author

MatzeB commented Dec 9, 2024

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

Copy link
Member

@dtcxzyw dtcxzyw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thank you!

@dtcxzyw dtcxzyw changed the title Match range check pattern with SExt [InstCombine] Match range check pattern with SExt Dec 10, 2024
@MatzeB MatzeB merged commit e9c68c6 into llvm:main Dec 10, 2024
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
llvm:instcombine Covers the InstCombine, InstSimplify and AggressiveInstCombine passes llvm:transforms
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants