[DAGCombiner] widen any_ext of popcount based on target support

This enhances D69127 (rGe6c145e0548e3b3de6eab27e44e1504387cf6b53)
to handle the looser "any_extend" cast in addition to zext.

This is a prerequisite step for canonicalizing in the other direction
(narrow the popcount) in IR - PR43688:
https://bugs.llvm.org/show_bug.cgi?id=43688

3 files changed, 36 insertions, 19 deletions

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 6b30ebecbea..5003c8f0536 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -9673,6 +9673,29 @@ static bool isTruncateOf(SelectionDAG &DAG, SDValue N, SDValue &Op,
   return (Known.Zero | 1).isAllOnesValue();
 }
 
+/// Given an extending node with a pop-count operand, if the target does not
+/// support a pop-count in the narrow source type but does support it in the
+/// destination type, widen the pop-count to the destination type.
+static SDValue widenCtPop(SDNode *Extend, SelectionDAG &DAG) {
+  assert((Extend->getOpcode() == ISD::ZERO_EXTEND ||
+          Extend->getOpcode() == ISD::ANY_EXTEND) && "Expected extend op");
+
+  SDValue CtPop = Extend->getOperand(0);
+  if (CtPop.getOpcode() != ISD::CTPOP || !CtPop.hasOneUse())
+    return SDValue();
+
+  EVT VT = Extend->getValueType(0);
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  if (TLI.isOperationLegalOrCustom(ISD::CTPOP, CtPop.getValueType()) ||
+      !TLI.isOperationLegalOrCustom(ISD::CTPOP, VT))
+    return SDValue();
+
+  // zext (ctpop X) --> ctpop (zext X)
+  SDLoc DL(Extend);
+  SDValue NewZext = DAG.getZExtOrTrunc(CtPop.getOperand(0), DL, VT);
+  return DAG.getNode(ISD::CTPOP, DL, VT, NewZext);
+}
+
 SDValue DAGCombiner::visitZERO_EXTEND(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
@@ -9923,17 +9946,8 @@ SDValue DAGCombiner::visitZERO_EXTEND(SDNode *N) {
   if (SDValue NewVSel = matchVSelectOpSizesWithSetCC(N))
     return NewVSel;
 
-  // If the target does not support a pop-count in the narrow source type but
-  // does support it in the destination type, widen the pop-count to this type:
-  // zext (ctpop X) --> ctpop (zext X)
-  // TODO: Generalize this to handle starting from anyext.
-  if (N0.getOpcode() == ISD::CTPOP && N0.hasOneUse() &&
-      !TLI.isOperationLegalOrCustom(ISD::CTPOP, N0.getValueType()) &&
-      TLI.isOperationLegalOrCustom(ISD::CTPOP, VT)) {
-    SDLoc DL(N);
-    SDValue NewZext = DAG.getZExtOrTrunc(N0.getOperand(0), DL, VT);
-    return DAG.getNode(ISD::CTPOP, DL, VT, NewZext);
-  }
+  if (SDValue NewCtPop = widenCtPop(N, DAG))
+    return NewCtPop;
 
   return SDValue();
 }
@@ -10081,6 +10095,9 @@ SDValue DAGCombiner::visitANY_EXTEND(SDNode *N) {
       return SCC;
   }
 
+  if (SDValue NewCtPop = widenCtPop(N, DAG))
+    return NewCtPop;
+
   return SDValue();
 }
 
diff --git a/llvm/test/CodeGen/PowerPC/popcnt-zext.ll b/llvm/test/CodeGen/PowerPC/popcnt-zext.ll
index 789e70b2c5b..23d49f2290f 100644
--- a/llvm/test/CodeGen/PowerPC/popcnt-zext.ll
+++ b/llvm/test/CodeGen/PowerPC/popcnt-zext.ll
@@ -5,8 +5,8 @@
 define i16 @zpop_i8_i16(i8 %x) {
 ; FAST-LABEL: zpop_i8_i16:
 ; FAST:       # %bb.0:
-; FAST-NEXT:    rlwinm 3, 3, 0, 24, 31
-; FAST-NEXT:    popcntw 3, 3
+; FAST-NEXT:    clrldi 3, 3, 56
+; FAST-NEXT:    popcntd 3, 3
 ; FAST-NEXT:    blr
 ;
 ; SLOW-LABEL: zpop_i8_i16:
@@ -297,8 +297,8 @@ define i64 @popz_i32_i64(i32 %x) {
 define i64 @popa_i16_i64(i16 %x) {
 ; FAST-LABEL: popa_i16_i64:
 ; FAST:       # %bb.0:
-; FAST-NEXT:    rlwinm 3, 3, 0, 16, 31
-; FAST-NEXT:    popcntw 3, 3
+; FAST-NEXT:    clrldi 3, 3, 48
+; FAST-NEXT:    popcntd 3, 3
 ; FAST-NEXT:    andi. 3, 3, 16
 ; FAST-NEXT:    blr
 ;
diff --git a/llvm/test/CodeGen/PowerPC/popcnt.ll b/llvm/test/CodeGen/PowerPC/popcnt.ll
index 9c71c94f510..11e1ec3ffe7 100644
--- a/llvm/test/CodeGen/PowerPC/popcnt.ll
+++ b/llvm/test/CodeGen/PowerPC/popcnt.ll
@@ -8,8 +8,8 @@ define i8 @cnt8(i8 %x) nounwind readnone {
   %cnt = tail call i8 @llvm.ctpop.i8(i8 %x)
   ret i8 %cnt
 ; CHECK-LABEL: @cnt8
-; CHECK: rlwinm
-; CHECK: popcntw
+; CHECK: clrldi
+; CHECK: popcntd
 ; CHECK: blr
 
 ; SLOWPC-LABEL: @cnt8
@@ -20,8 +20,8 @@ define i16 @cnt16(i16 %x) nounwind readnone {
   %cnt = tail call i16 @llvm.ctpop.i16(i16 %x)
   ret i16 %cnt
 ; CHECK-LABEL: @cnt16
-; CHECK: rlwinm
-; CHECK: popcntw
+; CHECK: clrldi
+; CHECK: popcntd
 ; CHECK: blr
 
 ; SLOWPC-LABEL: @cnt16