[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 0dddb816196eefe1d4fd98e3b9f7196d196a6306 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] [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]]
 ;