[llvm] [InstCombine] Fold `ctpop(X) <u 2` into `ctpop(X) == 1` if X is non-zero (PR #67268)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/67268

This patch folds pattern `ctpop(X) <u 2` into `ctpop(X) == 1` if we know `X` is non-zero.
Fixes #57328.

>From 511e83e2b6336dd647ed3e4dd4a5c2edc7512ce2 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sun, 24 Sep 2023 23:19:25 +0800
Subject: [PATCH] [InstCombine] Simplify pattern `isPow2OrZero` if X is
 non-zero

---
 llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp | 9 ++++++++-
 llvm/test/Transforms/InstCombine/ispow2.ll              | 2 +-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index a219dac7acfbe16..9aafd83d42d0756 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -3412,6 +3412,14 @@ static Instruction *foldCtpopPow2Test(ICmpInst &I, IntrinsicInst *CtpopLhs,
                                       const SimplifyQuery &Q) {
   assert(CtpopLhs->getIntrinsicID() == Intrinsic::ctpop &&
          "Non-ctpop intrin in ctpop fold");
+
+  const ICmpInst::Predicate Pred = I.getPredicate();
+  // If we know X is non-zero, we can fold isPow2OrZero into isPow2.
+  if (Pred == ICmpInst::ICMP_ULT && CRhs == 2 &&
+      isKnownNonZero(CtpopLhs, Q.DL, /*Depth*/ 0, Q.AC, Q.CxtI, Q.DT))
+    return ICmpInst::Create(Instruction::ICmp, ICmpInst::ICMP_EQ, CtpopLhs,
+                            ConstantInt::get(CtpopLhs->getType(), 1));
+
   if (!CtpopLhs->hasOneUse())
     return nullptr;
 
@@ -3423,7 +3431,6 @@ static Instruction *foldCtpopPow2Test(ICmpInst &I, IntrinsicInst *CtpopLhs,
   // If we know any bit of X can be folded to:
   //    IsPow2       : X & (~Bit) == 0
   //    NotPow2      : X & (~Bit) != 0
-  const ICmpInst::Predicate Pred = I.getPredicate();
   if (((I.isEquality() || Pred == ICmpInst::ICMP_UGT) && CRhs == 1) ||
       (Pred == ICmpInst::ICMP_ULT && CRhs == 2)) {
     Value *Op = CtpopLhs->getArgOperand(0);
diff --git a/llvm/test/Transforms/InstCombine/ispow2.ll b/llvm/test/Transforms/InstCombine/ispow2.ll
index bbd693b11b388ad..60eb522a144927f 100644
--- a/llvm/test/Transforms/InstCombine/ispow2.ll
+++ b/llvm/test/Transforms/InstCombine/ispow2.ll
@@ -198,7 +198,7 @@ define i1 @is_pow2_non_zero(i32 %x) {
 ; CHECK-NEXT:    [[NOTZERO:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[NOTZERO]])
 ; CHECK-NEXT:    [[T0:%.*]] = tail call i32 @llvm.ctpop.i32(i32 [[X]]), !range [[RNG0]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[T0]], 2
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[T0]], 1
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %notzero = icmp ne i32 %x, 0