[llvm] [ConstantRange] Handle `Intrinsic::ctpop` (PR #68310)

[Yingwei Zheng via llvm-commits]

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

This patch adds support for `Intrinsic::ctpop` in ConstantRange. It calculates the range in O(1) with the LCP-based method.

Migrated from https://reviews.llvm.org/D153505.


>From 0b1ce7278443d9f777d6af2fb022469b055158ad Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 5 Oct 2023 21:48:04 +0800
Subject: [PATCH] [ConstantRange] Handle `Intrinsic::ctpop`

---
 llvm/include/llvm/IR/ConstantRange.h          |  3 ++
 llvm/lib/IR/ConstantRange.cpp                 | 49 +++++++++++++++++++
 .../CorrelatedValuePropagation/range.ll       |  3 +-
 llvm/unittests/IR/ConstantRangeTest.cpp       |  6 +++
 4 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/IR/ConstantRange.h b/llvm/include/llvm/IR/ConstantRange.h
index ca36732e4e2e8c2..017f1f36d8a663e 100644
--- a/llvm/include/llvm/IR/ConstantRange.h
+++ b/llvm/include/llvm/IR/ConstantRange.h
@@ -530,6 +530,9 @@ class [[nodiscard]] ConstantRange {
   /// ignoring a possible zero value contained in the input range.
   ConstantRange ctlz(bool ZeroIsPoison = false) const;
 
+  /// Calculate ctpop range.
+  ConstantRange ctpop() const;
+
   /// Represents whether an operation on the given constant range is known to
   /// always or never overflow.
   enum class OverflowResult {
diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index 3d71b20f7e853e0..f00fb1f1b9f36cc 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -949,6 +949,7 @@ bool ConstantRange::isIntrinsicSupported(Intrinsic::ID IntrinsicID) {
   case Intrinsic::smax:
   case Intrinsic::abs:
   case Intrinsic::ctlz:
+  case Intrinsic::ctpop:
     return true;
   default:
     return false;
@@ -986,6 +987,8 @@ ConstantRange ConstantRange::intrinsic(Intrinsic::ID IntrinsicID,
     assert(ZeroIsPoison->getBitWidth() == 1 && "Must be boolean");
     return Ops[0].ctlz(ZeroIsPoison->getBoolValue());
   }
+  case Intrinsic::ctpop:
+    return Ops[0].ctpop();
   default:
     assert(!isIntrinsicSupported(IntrinsicID) && "Shouldn't be supported");
     llvm_unreachable("Unsupported intrinsic");
@@ -1736,6 +1739,52 @@ ConstantRange ConstantRange::ctlz(bool ZeroIsPoison) const {
                      APInt(getBitWidth(), getUnsignedMin().countl_zero() + 1));
 }
 
+static ConstantRange getUnsignedPopCountRange(const APInt &Lower,
+                                              const APInt &Upper) {
+  assert(Lower.ule(Upper) && "Unexpected wrapped set.");
+  unsigned BitWidth = Lower.getBitWidth();
+  if (Lower == Upper)
+    return ConstantRange::getEmpty(BitWidth);
+  if (Lower + 1 == Upper)
+    return ConstantRange(APInt(BitWidth, Lower.popcount()));
+
+  APInt Max = Upper - 1;
+  // Calculate longest common prefix.
+  unsigned LCPLength = (Lower ^ Max).countl_zero();
+  unsigned LCPPopCount = Lower.getHiBits(LCPLength).popcount();
+  // If Lower is {LCP, 000...}, the minimum is the popcount of LCP.
+  // Otherwise, the minimum is the popcount of LCP + 1.
+  unsigned MinBits =
+      LCPPopCount + (Lower.countr_zero() < BitWidth - LCPLength ? 1 : 0);
+  // If Max is {LCP, 111...}, the maximum is the popcount of LCP + (BitWidth -
+  // length of LCP).
+  // Otherwise, the minimum is the popcount of LCP + (BitWidth -
+  // length of LCP - 1).
+  unsigned MaxBits = LCPPopCount + (BitWidth - LCPLength) +
+                     (Max.countr_one() >= BitWidth - LCPLength ? 1 : 0);
+  return ConstantRange(APInt(BitWidth, MinBits), APInt(BitWidth, MaxBits));
+}
+
+ConstantRange ConstantRange::ctpop() const {
+  if (isEmptySet())
+    return getEmpty();
+
+  unsigned BitWidth = getBitWidth();
+  APInt Zero = APInt::getZero(BitWidth);
+  if (isFullSet()) {
+    return getNonEmpty(Zero, APInt(BitWidth, BitWidth + 1));
+  }
+  if (!isUpperWrapped()) {
+    return getUnsignedPopCountRange(getLower(), getUpper());
+  }
+  ConstantRange CR1 = ConstantRange(
+      APInt(BitWidth,
+            BitWidth - (getUnsignedMax() - getLower() + 1).logBase2()),
+      APInt(BitWidth, BitWidth + 1)); // [lower, intmax]
+  ConstantRange CR2 = getUnsignedPopCountRange(Zero, getUpper()); // [0, upper)
+  return CR1.unionWith(CR2);
+}
+
 ConstantRange::OverflowResult ConstantRange::unsignedAddMayOverflow(
     const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/range.ll b/llvm/test/Transforms/CorrelatedValuePropagation/range.ll
index 5c10fdcd2b49217..cf85753e59c8083 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/range.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/range.ll
@@ -1085,8 +1085,7 @@ define i1 @ctpop_fold(i16 %x) {
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    [[CTPOP:%.*]] = call i16 @llvm.ctpop.i16(i16 [[X]])
-; CHECK-NEXT:    [[RES:%.*]] = icmp ule i16 [[CTPOP]], 8
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ; CHECK:       else:
 ; CHECK-NEXT:    ret i1 true
 ;
diff --git a/llvm/unittests/IR/ConstantRangeTest.cpp b/llvm/unittests/IR/ConstantRangeTest.cpp
index 1cb358a26062ca5..12facfb22fb3c73 100644
--- a/llvm/unittests/IR/ConstantRangeTest.cpp
+++ b/llvm/unittests/IR/ConstantRangeTest.cpp
@@ -2438,6 +2438,12 @@ TEST_F(ConstantRangeTest, Ctlz) {
       });
 }
 
+TEST_F(ConstantRangeTest, Ctpop) {
+  TestUnaryOpExhaustive(
+      [](const ConstantRange &CR) { return CR.ctpop(); },
+      [](const APInt &N) { return APInt(N.getBitWidth(), N.popcount()); });
+}
+
 TEST_F(ConstantRangeTest, castOps) {
   ConstantRange A(APInt(16, 66), APInt(16, 128));
   ConstantRange FpToI8 = A.castOp(Instruction::FPToSI, 8);