[llvm] [InstCombine] Fold zext-of-icmp with no shift (PR #68503)

[via llvm-commits]

https://github.com/vfdff updated https://github.com/llvm/llvm-project/pull/68503

>From b4550f7f563287b86b7b99f1d4615782d040252d Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Sat, 7 Oct 2023 23:44:18 -0400
Subject: [PATCH 1/2] [InstCombine] Precommit test for PR68465

---
 llvm/test/Transforms/InstCombine/zext.ll | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/zext.ll b/llvm/test/Transforms/InstCombine/zext.ll
index 8aa2a10e6abb2ea..00147bb523f474e 100644
--- a/llvm/test/Transforms/InstCombine/zext.ll
+++ b/llvm/test/Transforms/InstCombine/zext.ll
@@ -748,3 +748,16 @@ define i64 @zext_icmp_ne_bool_1(ptr %ptr) {
   %len = zext i1 %cmp to i64
   ret i64 %len
 }
+
+define i32  @zext_icmp_eq0_no_shift(ptr %ptr ) {
+; CHECK-LABEL: @zext_icmp_eq0_no_shift(
+; CHECK-NEXT:    [[X:%.*]] = load i8, ptr [[PTR:%.*]], align 1, !range [[RNG1:![0-9]+]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[X]], 0
+; CHECK-NEXT:    [[RES:%.*]] = zext i8 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %X = load i8, ptr %ptr,align 1, !range !{i8 0, i8 2} ; range [0, 2)
+  %cmp = icmp eq i8 %X, 0
+  %res = zext i1 %cmp to i32
+  ret i32 %res
+}

>From 04dbbc020e788d1f4d6155819dc80b74b2520545 Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Sat, 7 Oct 2023 04:43:57 -0400
Subject: [PATCH 2/2] [InstCombine] Fold zext-of-icmp with no shift

This regression triggers after commit f400daa to fix infinite loop issue.

In this case, we can known the shift count is 0, so it will not be
triggered by the form of (iN (~X) u>> (N - 1)) in commit 21d3871, of
which N indicates the data type bitwidth of X.

Fixes https://github.com/llvm/llvm-project/issues/68465
---
 .../Transforms/InstCombine/InstCombineCasts.cpp | 17 ++++++++---------
 llvm/test/Transforms/InstCombine/zext.ll        |  3 ++-
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 20c13de33f8189d..5e23514bf5ae0be 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -904,19 +904,18 @@ Instruction *InstCombinerImpl::transformZExtICmp(ICmpInst *Cmp,
     // zext (X == 0) to i32 --> (X>>1)^1 iff X has only the 2nd bit set.
     // zext (X != 0) to i32 --> X        iff X has only the low bit set.
     // zext (X != 0) to i32 --> X>>1     iff X has only the 2nd bit set.
-    if (Op1CV->isZero() && Cmp->isEquality() &&
-        (Cmp->getOperand(0)->getType() == Zext.getType() ||
-         Cmp->getPredicate() == ICmpInst::ICMP_NE)) {
-      // If Op1C some other power of two, convert:
-      KnownBits Known = computeKnownBits(Cmp->getOperand(0), 0, &Zext);
 
+    if (Op1CV->isZero() && Cmp->isEquality()) {
       // Exactly 1 possible 1? But not the high-bit because that is
       // canonicalized to this form.
+      KnownBits Known = computeKnownBits(Cmp->getOperand(0), 0, &Zext);
       APInt KnownZeroMask(~Known.Zero);
-      if (KnownZeroMask.isPowerOf2() &&
-          (Zext.getType()->getScalarSizeInBits() !=
-           KnownZeroMask.logBase2() + 1)) {
-        uint32_t ShAmt = KnownZeroMask.logBase2();
+      uint32_t ShAmt = KnownZeroMask.logBase2();
+      bool IsExpectShAmt = KnownZeroMask.isPowerOf2() &&
+                           (Zext.getType()->getScalarSizeInBits() != ShAmt + 1);
+      if (IsExpectShAmt &&
+          (Cmp->getOperand(0)->getType() == Zext.getType() ||
+           Cmp->getPredicate() == ICmpInst::ICMP_NE || ShAmt == 0)) {
         Value *In = Cmp->getOperand(0);
         if (ShAmt) {
           // Perform a logical shr by shiftamt.
diff --git a/llvm/test/Transforms/InstCombine/zext.ll b/llvm/test/Transforms/InstCombine/zext.ll
index 00147bb523f474e..f595638ba9e3083 100644
--- a/llvm/test/Transforms/InstCombine/zext.ll
+++ b/llvm/test/Transforms/InstCombine/zext.ll
@@ -749,10 +749,11 @@ define i64 @zext_icmp_ne_bool_1(ptr %ptr) {
   ret i64 %len
 }
 
+; https://alive2.llvm.org/ce/z/k7qosS
 define i32  @zext_icmp_eq0_no_shift(ptr %ptr ) {
 ; CHECK-LABEL: @zext_icmp_eq0_no_shift(
 ; CHECK-NEXT:    [[X:%.*]] = load i8, ptr [[PTR:%.*]], align 1, !range [[RNG1:![0-9]+]]
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 [[X]], 0
+; CHECK-NEXT:    [[TMP1:%.*]] = xor i8 [[X]], 1
 ; CHECK-NEXT:    [[RES:%.*]] = zext i8 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[RES]]
 ;