[DAGCombiner] widen zext of popcount based on target support

[spatel]

Summary

zext (ctpop X) --> ctpop (zext X)

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

I'm not sure if any other targets are affected, but I found a missing fold for PPC, so added tests based on that.
The reason we widen all the way to 64-bit in these tests is because the initial DAG looks something like this:

t5: i8 = ctpop t4
t6: i32 = zero_extend t5  <-- created based on IR, but unused node?
  t7: i64 = zero_extend t5

[xbolva00]

Looks good

[jsji]

LGTM. Thanks.
Do we want to consider any_extend as well?

[spatel]

In D69127#1715078, @jsji wrote:

LGTM. Thanks.
Do we want to consider any_extend as well?

I guess that's possible (convert the anyext to zext?) although I'm not sure how to create a test case.

For example, if we had:
i16 = anyext (i8 = ctpop 0xF0)

The result must be "i16 0x??04" because we only know the value of the lower bits exactly.

i16 = ctpop (i16 = zext 0xF0)

The result must be "i16 0x0004", so it's more defined than the original, but it's probably still a perf win to use a supported ctpop operation?

[jsji]

Yeah.

I am asking because I see we will use any_extend with the simple call.

define i8 @popz_i8_i8(i8 %x) {
  %pop = tail call i8 @llvm.ctpop.i8(i8 %x)
  ret i8 %pop
}

widening in this example might not get a perf gain,
but I think in some case we might win, at least not loss?

And since you mention this is a prerequisite for canonicalizing,
so maybe that will bring more possibility of win?

[craig.topper]

So the assumption here is that the popcnt we choose here is the same performance as any smaller popcnt that LegalizeDAG would pick through promotion. Or near enough that its better than having 2 zexts?

[RKSimon]

any_extend might get generated when using ctpop for 'allones' tests:

define i64 @var_ctpop_i32(i32 %a) {
  %1 = call i32 @llvm.ctpop.i32(i32 %a)
  %2 = zext i32 %1 to i64 ; SimplifyDemandedBits may turn zext (or sext) into aext
  %3 = and i64 %2, 32
  ret i64 %3
}
declare i32 @llvm.ctpop.i32(i32)

[spatel]

In D69127#1715373, @craig.topper wrote:

So the assumption here is that the popcnt we choose here is the same performance as any smaller popcnt that LegalizeDAG would pick through promotion. Or near enough that its better than having 2 zexts?

Almost - we're saying that ctpop+zext is no better perf than the wider ctpop. For example, the i8 test on PPC ends up like this after promotion without this patch:

    t13: i32 = and t11, Constant:i32<255>
  t14: i32 = ctpop t13
t18: i64 = zero_extend t14

This patch would already have hoisted the zext ahead of the i32 ctpop in that basic example, but if we had started with this i32 code, then we'd widen to i64 ctpop because we assume it's no more expensive than i32 ctpop+zext. For the PPC examples, the zext becomes a mask op, so we eliminate that instruction.

We shouldn't be touching anything that's expanded. So for base x86, we're not going to widen an i8 op because that expansion should be cheaper than any wider expansion.

[spatel]

In D69127#1715640, @spatel wrote:

This patch would already have hoisted the zext ahead of the i32 ctpop in that basic example, but if we had started with this i32 code, then we'd widen to i64 ctpop because we assume it's no more expensive than i32 ctpop+zext. For the PPC examples, the zext becomes a mask op, so we eliminate that instruction.

Oops - that last part was wrong. (And there's already a test for i32 -> i64.) Since PPC has legal ops for both types, we won't do anything for that pattern.

[spatel]

In D69127#1715638, @RKSimon wrote:

any_extend might get generated when using ctpop for 'allones' tests:

define i64 @var_ctpop_i32(i32 %a) {
  %1 = call i32 @llvm.ctpop.i32(i32 %a)
  %2 = zext i32 %1 to i64 ; SimplifyDemandedBits may turn zext (or sext) into aext
  %3 = and i64 %2, 32
  ret i64 %3
}
declare i32 @llvm.ctpop.i32(i32)

Thanks - I added a variant of this in:
rGb74d7e5cccb5
I'll add a TODO comment here and follow-up to generalize for that case.

[xbolva00]

lg

[jsji]

I'll add a TODO comment here and follow-up to generalize for that case.

Great. Thanks @spatel !

[spatel]

any_ext added with:
rG1ebd4a2e3ad0

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if ((N0.getOpcode() == ISD::SHL || N0.getOpcode() == ISD::SRL) &&
     return DAG.getNode(N0.getOpcode(), DL, VT,
                        DAG.getNode(ISD::ZERO_EXTEND, DL, VT, N0.getOperand(0)),
                        ShAmt);
   }
 
   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);
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitANY_EXTEND(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
   if (SDValue Res = tryToFoldExtendOfConstant(N, TLI, DAG, LegalTypes))

llvm/test/CodeGen/PowerPC/popcnt-zext.ll

; SLOW-NEXT:    blr
   %z = zext i8 %x to i16
   %pop = tail call i16 @llvm.ctpop.i16(i16 %z)
   ret i16 %pop
 }
 
 define i16 @popz_i8_i16(i8 %x) {
 ; FAST-LABEL: popz_i8_i16:
 ; FAST:       # %bb.0:
-; FAST-NEXT:    rlwinm 3, 3, 0, 24, 31
-; FAST-NEXT:    popcntw 3, 3
-; FAST-NEXT:    clrldi 3, 3, 32
+; FAST-NEXT:    clrldi 3, 3, 56
+; FAST-NEXT:    popcntd 3, 3
 ; FAST-NEXT:    blr
 ;
 ; SLOW-LABEL: popz_i8_i16:
 ; SLOW:       # %bb.0:
 ; SLOW-NEXT:    clrlwi 5, 3, 24
 ; SLOW-NEXT:    rlwinm 3, 3, 31, 0, 31
 ; SLOW-NEXT:    andi. 3, 3, 85
 ; SLOW-NEXT:    lis 4, 13107
; SLOW-NEXT:    blr
   %z = zext i8 %x to i32
   %pop = tail call i32 @llvm.ctpop.i32(i32 %z)
   ret i32 %pop
 }
 
 define i32 @popz_i8_32(i8 %x) {
 ; FAST-LABEL: popz_i8_32:
 ; FAST:       # %bb.0:
-; FAST-NEXT:    rlwinm 3, 3, 0, 24, 31
-; FAST-NEXT:    popcntw 3, 3
-; FAST-NEXT:    clrldi 3, 3, 32
+; FAST-NEXT:    clrldi 3, 3, 56
+; FAST-NEXT:    popcntd 3, 3
 ; FAST-NEXT:    blr
 ;
 ; SLOW-LABEL: popz_i8_32:
 ; SLOW:       # %bb.0:
 ; SLOW-NEXT:    clrlwi 5, 3, 24
 ; SLOW-NEXT:    rlwinm 3, 3, 31, 0, 31
 ; SLOW-NEXT:    andi. 3, 3, 85
 ; SLOW-NEXT:    lis 4, 13107
; SLOW-NEXT:    blr
   %z = zext i16 %x to i32
   %pop = tail call i32 @llvm.ctpop.i32(i32 %z)
   ret i32 %pop
 }
 
 define i32 @popz_i16_32(i16 %x) {
 ; FAST-LABEL: popz_i16_32:
 ; FAST:       # %bb.0:
-; FAST-NEXT:    rlwinm 3, 3, 0, 16, 31
-; FAST-NEXT:    popcntw 3, 3
-; FAST-NEXT:    clrldi 3, 3, 32
+; FAST-NEXT:    clrldi 3, 3, 48
+; FAST-NEXT:    popcntd 3, 3
 ; FAST-NEXT:    blr
 ;
 ; SLOW-LABEL: popz_i16_32:
 ; SLOW:       # %bb.0:
 ; SLOW-NEXT:    clrlwi 5, 3, 16
 ; SLOW-NEXT:    rlwinm 3, 3, 31, 0, 31
 ; SLOW-NEXT:    andi. 3, 3, 21845
 ; SLOW-NEXT:    lis 4, 13107