[llvm] [WIP][DAG] Improve the knownbits of CTPOP/LZ/TZ based off the min/max counts of the input value (PR #107085)

[Simon Pilgrim via llvm-commits]

https://github.com/RKSimon updated https://github.com/llvm/llvm-project/pull/107085

>From b16401767b14d4718322e19d453e6d5a91e97101 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Tue, 3 Sep 2024 12:06:45 +0100
Subject: [PATCH 1/2] [WIP][DAG] Improve the knownbits of CTPOP/LZ/TZ based off
 the min/max counts of the input value

Still need to add proper test coverage and decide whether we want to use ConstantRange or not
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 20 +++++++++----------
 llvm/test/CodeGen/AArch64/pr61549.ll          |  8 +-------
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index aa468fa9ebb4c3..d4c633bf017adb 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3747,26 +3747,24 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
   case ISD::CTTZ:
   case ISD::CTTZ_ZERO_UNDEF: {
     Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
-    // If we have a known 1, its position is our upper bound.
-    unsigned PossibleTZ = Known2.countMaxTrailingZeros();
-    unsigned LowBits = llvm::bit_width(PossibleTZ);
-    Known.Zero.setBitsFrom(LowBits);
+    APInt LowerCount(BitWidth, Known2.countMinTrailingZeros());
+    APInt UpperCount(BitWidth, 1 + Known2.countMaxTrailingZeros());
+    Known = ConstantRange(LowerCount, UpperCount).toKnownBits();
     break;
   }
   case ISD::CTLZ:
   case ISD::CTLZ_ZERO_UNDEF: {
     Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
-    // If we have a known 1, its position is our upper bound.
-    unsigned PossibleLZ = Known2.countMaxLeadingZeros();
-    unsigned LowBits = llvm::bit_width(PossibleLZ);
-    Known.Zero.setBitsFrom(LowBits);
+    APInt LowerCount(BitWidth, Known2.countMinLeadingZeros());
+    APInt UpperCount(BitWidth, 1 + Known2.countMaxLeadingZeros());
+    Known = ConstantRange(LowerCount, UpperCount).toKnownBits();
     break;
   }
   case ISD::CTPOP: {
     Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
-    // If we know some of the bits are zero, they can't be one.
-    unsigned PossibleOnes = Known2.countMaxPopulation();
-    Known.Zero.setBitsFrom(llvm::bit_width(PossibleOnes));
+    APInt LowerCount(BitWidth, Known2.countMinPopulation());
+    APInt UpperCount(BitWidth, 1 + Known2.countMaxPopulation());
+    Known = ConstantRange(LowerCount, UpperCount).toKnownBits();
     break;
   }
   case ISD::PARITY: {
diff --git a/llvm/test/CodeGen/AArch64/pr61549.ll b/llvm/test/CodeGen/AArch64/pr61549.ll
index e66ee7d219cc5e..b392c418cd9060 100644
--- a/llvm/test/CodeGen/AArch64/pr61549.ll
+++ b/llvm/test/CodeGen/AArch64/pr61549.ll
@@ -5,13 +5,7 @@
 define i35 @f(i35 %0) {
 ; CHECK-LABEL: f:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sbfx x8, x0, #0, #35
-; CHECK-NEXT:    mov w9, #1 // =0x1
-; CHECK-NEXT:    sdiv x10, x9, x8
-; CHECK-NEXT:    msub x8, x10, x8, x9
-; CHECK-NEXT:    clz x8, x8
-; CHECK-NEXT:    sub x8, x8, #29
-; CHECK-NEXT:    ubfx x0, x8, #5, #30
+; CHECK-NEXT:    mov w0, #1 // =0x1
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: f:

>From b51768c1d7b0ae87068bfc9a83bf261b03646224 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev at redking.me.uk>
Date: Tue, 3 Sep 2024 13:31:53 +0100
Subject: [PATCH 2/2] Use ConstantRange::getNonEmpty

---
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index d4c633bf017adb..2e0d7829c933a8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3747,24 +3747,24 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
   case ISD::CTTZ:
   case ISD::CTTZ_ZERO_UNDEF: {
     Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
-    APInt LowerCount(BitWidth, Known2.countMinTrailingZeros());
-    APInt UpperCount(BitWidth, 1 + Known2.countMaxTrailingZeros());
-    Known = ConstantRange(LowerCount, UpperCount).toKnownBits();
+    APInt MinCount(BitWidth, Known2.countMinTrailingZeros());
+    APInt MaxCount(BitWidth, Known2.countMaxTrailingZeros());
+    Known = ConstantRange::getNonEmpty(MinCount, MaxCount + 1).toKnownBits();
     break;
   }
   case ISD::CTLZ:
   case ISD::CTLZ_ZERO_UNDEF: {
     Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
-    APInt LowerCount(BitWidth, Known2.countMinLeadingZeros());
-    APInt UpperCount(BitWidth, 1 + Known2.countMaxLeadingZeros());
-    Known = ConstantRange(LowerCount, UpperCount).toKnownBits();
+    APInt MinCount(BitWidth, Known2.countMinLeadingZeros());
+    APInt MaxCount(BitWidth, Known2.countMaxLeadingZeros());
+    Known = ConstantRange::getNonEmpty(MinCount, MaxCount + 1).toKnownBits();
     break;
   }
   case ISD::CTPOP: {
     Known2 = computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
-    APInt LowerCount(BitWidth, Known2.countMinPopulation());
-    APInt UpperCount(BitWidth, 1 + Known2.countMaxPopulation());
-    Known = ConstantRange(LowerCount, UpperCount).toKnownBits();
+    APInt MinCount(BitWidth, Known2.countMinPopulation());
+    APInt MaxCount(BitWidth, Known2.countMaxPopulation());
+    Known = ConstantRange::getNonEmpty(MinCount, MaxCount + 1).toKnownBits();
     break;
   }
   case ISD::PARITY: {