[llvm] b279ca2 - [DAG] visitCTPOP - CTPOP(SHIFT(X)) -> CTPOP(X) iff the shift doesn't affect any non-zero bits

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2024-02-15T10:41:08Z
New Revision: b279ca278370dabd27b8d380d59bdb5018366053

URL: https://github.com/llvm/llvm-project/commit/b279ca278370dabd27b8d380d59bdb5018366053
DIFF: https://github.com/llvm/llvm-project/commit/b279ca278370dabd27b8d380d59bdb5018366053.diff

LOG: [DAG] visitCTPOP - CTPOP(SHIFT(X)) -> CTPOP(X) iff the shift doesn't affect any non-zero bits

If the source is being (logically) shifted, but doesn't affect any active bits, then we can call CTPOP on the shift source directly.

Added: 
    

Modified: 
    llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
    llvm/test/CodeGen/X86/ctpop-mask.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index bdd2336fa42379..18037f11ee4088 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -11144,6 +11144,23 @@ SDValue DAGCombiner::visitCTPOP(SDNode *N) {
   if (SDValue C = DAG.FoldConstantArithmetic(ISD::CTPOP, DL, VT, {N0}))
     return C;
 
+  // If the source is being shifted, but doesn't affect any active bits,
+  // then we can call CTPOP on the shift source directly.
+  if (N0.getOpcode() == ISD::SRL || N0.getOpcode() == ISD::SHL) {
+    if (ConstantSDNode *AmtC = isConstOrConstSplat(N0.getOperand(1))) {
+      const APInt &Amt = AmtC->getAPIntValue();
+      if (Amt.ult(NumBits)) {
+        KnownBits KnownSrc = DAG.computeKnownBits(N0.getOperand(0));
+        if ((N0.getOpcode() == ISD::SRL &&
+             Amt.ule(KnownSrc.countMinTrailingZeros())) ||
+            (N0.getOpcode() == ISD::SHL &&
+             Amt.ule(KnownSrc.countMinLeadingZeros()))) {
+          return DAG.getNode(ISD::CTPOP, DL, VT, N0.getOperand(0));
+        }
+      }
+    }
+  }
+
   // If the upper bits are known to be zero, then see if its profitable to
   // only count the lower bits.
   if (VT.isScalarInteger() && NumBits > 8 && (NumBits & 1) == 0) {

diff  --git a/llvm/test/CodeGen/X86/ctpop-mask.ll b/llvm/test/CodeGen/X86/ctpop-mask.ll
index 97c634a2a133d7..602d9b511cdc06 100644
--- a/llvm/test/CodeGen/X86/ctpop-mask.ll
+++ b/llvm/test/CodeGen/X86/ctpop-mask.ll
@@ -549,7 +549,6 @@ define i64 @ctpop_shifted_mask8(i64 %x) nounwind readnone {
 ; X86-POPCOUNT-LABEL: ctpop_shifted_mask8:
 ; X86-POPCOUNT:       # %bb.0:
 ; X86-POPCOUNT-NEXT:    movzbl {{[0-9]+}}(%esp), %eax
-; X86-POPCOUNT-NEXT:    shll $8, %eax
 ; X86-POPCOUNT-NEXT:    popcntl %eax, %eax
 ; X86-POPCOUNT-NEXT:    xorl %edx, %edx
 ; X86-POPCOUNT-NEXT:    retl
@@ -663,12 +662,11 @@ define i64 @ctpop_shifted_mask16(i64 %x) nounwind readnone {
 ; X86-NO-POPCOUNT-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NO-POPCOUNT-NEXT:    movl %ecx, %eax
 ; X86-NO-POPCOUNT-NEXT:    andl $524280, %eax # imm = 0x7FFF8
-; X86-NO-POPCOUNT-NEXT:    shrl $4, %ecx
-; X86-NO-POPCOUNT-NEXT:    andl $21845, %ecx # imm = 0x5555
-; X86-NO-POPCOUNT-NEXT:    shrl $3, %eax
+; X86-NO-POPCOUNT-NEXT:    shrl %ecx
+; X86-NO-POPCOUNT-NEXT:    andl $87380, %ecx # imm = 0x15554
 ; X86-NO-POPCOUNT-NEXT:    subl %ecx, %eax
 ; X86-NO-POPCOUNT-NEXT:    movl %eax, %ecx
-; X86-NO-POPCOUNT-NEXT:    andl $858993459, %ecx # imm = 0x33333333
+; X86-NO-POPCOUNT-NEXT:    andl $858993456, %ecx # imm = 0x33333330
 ; X86-NO-POPCOUNT-NEXT:    shrl $2, %eax
 ; X86-NO-POPCOUNT-NEXT:    andl $858993459, %eax # imm = 0x33333333
 ; X86-NO-POPCOUNT-NEXT:    addl %ecx, %eax