[llvm] [InstCombine] Fold zext(icmp ne X, C) to X - C when X is in (C, C+1) (PR #208914)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: muhammadshayanasim1-web

<details>
<summary>Changes</summary>

Addresses #<!-- -->207218

When the value of `X` is in `(CST, CST + 1)`, `zext(X != CST)` 
is equivalent to `X - CST`. This fold removes the icmp/zext pair, 
and replaces it with a single sub.

This PR only implments the first case from the issue (`src_zext_ne_lower_cst`), while 
the remaining three are left as follow-up.

Added test coverage for scalar/vector cases, widening zext, multi-use 
icmp, wrapped ranges, and negative cases.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp (+22) 
- (modified) llvm/test/Transforms/InstCombine/zext.ll (+93) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index f30a7b5e891ab..bf7fda20f012c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -1407,6 +1407,28 @@ Instruction *InstCombinerImpl::transformZExtICmp(ICmpInst *Cmp,
         return replaceInstUsesWith(Zext, IntCast);
       }
     }
+
+    ConstantRange CR = computeConstantRange(Cmp->getOperand(0), false, SQ);
+
+    if (Cmp->isEquality() && !CR.isSizeLargerThan(2) && CR.contains(*Op1CV)) {
+
+      Value *In = Cmp->getOperand(0);
+      bool SkipTransform = In->getType() != Zext.getType() && !Cmp->hasOneUse();
+
+      if (CR.getLower().eq(*Op1CV) &&
+          Cmp->getPredicate() == ICmpInst::ICMP_NE && !SkipTransform) {
+
+        In = Builder.CreateSub(In, ConstantInt::get(In->getType(), *Op1CV));
+
+        if (In->getType() == Zext.getType())
+          return replaceInstUsesWith(Zext, In);
+
+        if (Cmp->hasOneUse()) {
+          Value *IntCast = Builder.CreateIntCast(In, Zext.getType(), false);
+          return replaceInstUsesWith(Zext, IntCast);
+        }
+      }
+    }
   }
 
   if (Cmp->isEquality()) {
diff --git a/llvm/test/Transforms/InstCombine/zext.ll b/llvm/test/Transforms/InstCombine/zext.ll
index 3fb1e77ae2335..7faab734c6623 100644
--- a/llvm/test/Transforms/InstCombine/zext.ll
+++ b/llvm/test/Transforms/InstCombine/zext.ll
@@ -1080,3 +1080,96 @@ define <2 x i8> @zext_or_trunc_nuw_vec(<2 x i8> %x, <2 x i4> %y) {
   %zext = zext <2 x i4> %or to <2 x i8>
   ret <2 x i8> %zext
 }
+
+; zext (X != CST) --> X - CST        if X in {CST, CST + 1}
+
+define i32 @zext_ne_lower_cst_sub(i32 range(i32 7, 9) %x) {
+; CHECK-LABEL: @zext_ne_lower_cst_sub(
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[X:%.*]], -7
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, 7
+  %2 = zext i1 %1 to i32
+  ret i32 %2
+}
+
+define <4 x i32> @zext_ne_lower_cst_sub_vec(<4 x i32> range(i32 7, 9) %x) {
+; CHECK-LABEL: @zext_ne_lower_cst_sub_vec(
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw <4 x i32> [[X:%.*]], splat (i32 -7)
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
+;
+  %1 = icmp ne <4 x i32> %x, splat(i32 7)
+  %2 = zext <4 x i1> %1 to <4 x i32>
+  ret <4 x i32> %2
+}
+
+define i64 @zext_ne_lower_cst_sub_widen(i32 range(i32 7, 9) %x) {
+; CHECK-LABEL: @zext_ne_lower_cst_sub_widen(
+; CHECK-NEXT:    [[TMP1:%.*]] = add nsw i32 [[X:%.*]], -7
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i32 [[TMP1]] to i64
+; CHECK-NEXT:    ret i64 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, 7
+  %2 = zext i1 %1 to i64
+  ret i64 %2
+}
+
+define i32 @zext_ne_multiuse_icmp(i32 range(i32 7, 9) %x) {
+; CHECK-LABEL: @zext_ne_multiuse_icmp(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 7
+; CHECK-NEXT:    call void @use1(i1 [[TMP1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[X]], -7
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, 7
+  call void @use1(i1 %1)
+  %2 = zext i1 %1 to i32
+  ret i32 %2
+}
+
+define i32 @zext_ne_lower_wrapped(i32 range(i32 -1, 1) %x) {
+; CHECK-LABEL: @zext_ne_lower_wrapped(
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[X:%.*]], 1
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, -1
+  %2 = zext i1 %1 to i32
+  ret i32 %2
+}
+
+; Negative cases
+
+define i32 @neg_zext_ne_range_too_wide(i32 range(i32 7, 10) %x) {
+; CHECK-LABEL: @neg_zext_ne_range_too_wide(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 7
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i1 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, 7
+  %2 = zext i1 %1 to i32
+  ret i32 %2
+}
+
+define i32 @neg_zext_ne_no_range(i32 %x) {
+; CHECK-LABEL: @neg_zext_ne_no_range(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 7
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i1 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, 7
+  %2 = zext i1 %1 to i32
+  ret i32 %2
+}
+
+define i64 @neg_zext_ne_widen_multiuse_icmp(i32 range(i32 7, 9) %x) {
+; CHECK-LABEL: @neg_zext_ne_widen_multiuse_icmp(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 7
+; CHECK-NEXT:    call void @use1(i1 [[TMP1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i1 [[TMP1]] to i64
+; CHECK-NEXT:    ret i64 [[TMP2]]
+;
+  %1 = icmp ne i32 %x, 7
+  call void @use1(i1 %1)
+  %2 = zext i1 %1 to i64
+  ret i64 %2
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/208914