[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.

Author: RKSimon

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) {

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