[InstCombine] Match range check pattern with SExt (llvm#118910)

= 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
iree-org · Dec 10, 2024 · e9c68c6 · e9c68c6
1 parent fd57946
commit e9c68c6
Show file tree

Hide file tree

Showing 2 changed files with 144 additions and 4 deletions.
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -695,13 +695,17 @@ 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 (match(Cmp1Op0, m_SExtOrSelf(m_Specific(Input)))) {
     // For the upper range compare we have: icmp x, n
-    RangeEnd = Cmp1->getOperand(1);
-  } else if (Cmp1->getOperand(1) == Input) {
+    Input = Cmp1Op0;
+    RangeEnd = Cmp1Op1;
+  } else if (match(Cmp1Op1, m_SExtOrSelf(m_Specific(Input)))) {
     // For the upper range compare we have: icmp n, x
-    RangeEnd = Cmp1->getOperand(0);
+    Input = Cmp1Op1;
+    RangeEnd = Cmp1Op0;
     Pred1 = ICmpInst::getSwappedPredicate(Pred1);
   } else {
     return nullptr;

diff --git a/llvm/test/Transforms/InstCombine/range-check.ll b/llvm/test/Transforms/InstCombine/range-check.ll
@@ -32,6 +32,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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
+  ret i1 %c
+}
+
 define i1 @test_and2(i32 %x, i32 %n) {
 ; CHECK-LABEL: @test_and2(
 ; CHECK-NEXT:    [[NN:%.*]] = and i32 [[N:%.*]], 2147483647
@@ -60,6 +77,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp uge i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 sgt i32 %x, -1
+  %b = icmp sle i64 %x_sext, %n
+  %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 +120,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 sgt i64 %n, %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 +163,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp uge i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 i64 %n, %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 +208,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 slt i32 %x, 0
+  %b = icmp sge i64 %x_sext, %n
+  %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 +253,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 sle i32 %x, -1
+  %b = icmp sgt i64 %x_sext, %n
+  %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 +296,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ule i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 sle i64 %n, %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 +339,23 @@ 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:    [[N_NOT_NEGATIVE:%.*]] = icmp sgt i64 [[NN:%.*]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[N_NOT_NEGATIVE]])
+; CHECK-NEXT:    [[X_SEXT:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[C:%.*]] = icmp ult i64 [[NN]], [[X_SEXT]]
+; CHECK-NEXT:    ret i1 [[C]]
+;
+  %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 slt i64 %n, %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) {