[SDAG] expand is-power-of-2 pattern that uses popcount

[spatel]

Summary

(ctpop x) == 1 --> (x != 0) && ((x & x-1) == 0)

Adjust the legality check to avoid the poor codegen on AArch64. We probably only want to use popcount on this pattern when it is a single instruction.

This should fix issue #57225

[davezarzycki]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

18593–18594 ↗

(On Diff #454006)

For your consideration:

(ctpop x) <= 1 --> (x & x-1) == 0
(ctpop x) > 1 --> (x & x-1) != 0

[spatel]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

18593–18594 ↗

(On Diff #454006)

Ah, I forgot that we added a popcount target hook with D89952. Do you think we should use that for this pattern too? Or did you decide it wasn't worth pursuing (D91093)?

I'm looking back at the comments in:
https://bugs.llvm.org/show_bug.cgi?id=47825
...and we cited AArch64 scalar custom lowering/combining as a missing piece. But we could adapt the TLI hook if that seems better.

[davezarzycki]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

18593–18594 ↗

(On Diff #454006)

I think the D91093 work is good, but it stopped being a good use of my time to pursue. (Sorry)

Nevertheless the target hook is there. This seems like a reasonable use of it, no?

[spatel]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

18593–18594 ↗

(On Diff #454006)

Reviewing the earlier versions of D89952 - it doesn't seem like this pattern really fits the intent of the TLI hook?

This pattern will always expand to exactly one subtract+mask (+null check), so it's always going to be cheap based on the cost metric.

But we also have a legality check preceding both of these blocks. We can clean this up a bit and still have a smaller patch that gets us the expected results, so I'll post that alternative and see if it looks better.

[spatel]

Patch updated:
Adjust the legality checks instead of adding custom combining for AArch64.
This preserves x86 codegen for these patterns by keeping the hacky vector type check in place (without that we'd create a likely slower scalar "popcnt+cmp+set" instead of an "lea+test+set".

[dmgreen]

I had almost written something into 57225 that said, effectively, custom isn't a very good costmodel as it can mean so many things, so adding a target hook sounds like a good way forward.

This sounds good for AArch64. Could there be cases where the CTPOP (+setcc) is custom lowered but cheaper than the Add+And+setcc+setcc+and/or? My guess would probably be no, it would either be legal or quite expensive.

[spatel]

In D132237#3741964, @dmgreen wrote:

I had almost written something into 57225 that said, effectively, custom isn't a very good costmodel as it can mean so many things, so adding a target hook sounds like a good way forward.

This sounds good for AArch64. Could there be cases where the CTPOP (+setcc) is custom lowered but cheaper than the Add+And+setcc+setcc+and/or? My guess would probably be no, it would either be legal or quite expensive.

It's a fuzzy/general problem - we generally assume that a legal op is as cheap as it gets (usually one instruction), and a custom op is not as cheap (some sequence of legal instructions), but there's no guarantee on that.

I haven't seen any examples where this pattern is going wrong yet (after this patch), but it is complicated - @davezarzycki was looking at a large/messy set of potential x86 vector ISA trade-offs.

As noted earlier, the existing TLI hook doesn't quite fit the usage we want here. We could alter it to return a cost of 0 or something like that, but I thought that muddied the definition for the intended pattern (where we're expanding in a loop).

[dmgreen]

I agree, this looks like a good improvement. It might be worth adding a i32 version of the popcount compare tests. Otherwise LGTM

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

if (N0.getOpcode() == ISD::TRUNCATE && N0.hasOneUse() && !VT.isVector() &&
     CTPOP = N0.getOperand(0);
 
   if (CTPOP.getOpcode() != ISD::CTPOP || !CTPOP.hasOneUse())
     return SDValue();
 
   EVT CTVT = CTPOP.getValueType();
   SDValue CTOp = CTPOP.getOperand(0);
 
-  // If this is a vector CTPOP, keep the CTPOP if it is legal.
-  // TODO: Should we check if CTPOP is legal(or custom) for scalars?
-  if (VT.isVector() && TLI.isOperationLegal(ISD::CTPOP, CTVT))
-    return SDValue();
-
+  // Expand a power-of-2-or-zero comparison based on ctpop:
   // (ctpop x) u< 2 -> (x & x-1) == 0
   // (ctpop x) u> 1 -> (x & x-1) != 0
   if (Cond == ISD::SETULT || Cond == ISD::SETUGT) {
+    // Keep the CTPOP if it is a legal vector op.
+    if (CTVT.isVector() && TLI.isOperationLegal(ISD::CTPOP, CTVT))
+      return SDValue();
+
     unsigned CostLimit = TLI.getCustomCtpopCost(CTVT, Cond);
     if (C1.ugt(CostLimit + (Cond == ISD::SETULT)))
       return SDValue();
     if (C1 == 0 && (Cond == ISD::SETULT))
       return SDValue(); // This is handled elsewhere.
 
     unsigned Passes = C1.getLimitedValue() - (Cond == ISD::SETULT);
 
     SDValue NegOne = DAG.getAllOnesConstant(dl, CTVT);
     SDValue Result = CTOp;
     for (unsigned i = 0; i < Passes; i++) {
       SDValue Add = DAG.getNode(ISD::ADD, dl, CTVT, Result, NegOne);
       Result = DAG.getNode(ISD::AND, dl, CTVT, Result, Add);
     }
     ISD::CondCode CC = Cond == ISD::SETULT ? ISD::SETEQ : ISD::SETNE;
     return DAG.getSetCC(dl, VT, Result, DAG.getConstant(0, dl, CTVT), CC);
   }
 
-  // If ctpop is not supported, expand a power-of-2 comparison based on it.
+  // Expand a power-of-2 comparison based on ctpop:
+  // (ctpop x) == 1 --> (x != 0) && ((x & x-1) == 0)
+  // (ctpop x) != 1 --> (x == 0) || ((x & x-1) != 0)
   if ((Cond == ISD::SETEQ || Cond == ISD::SETNE) && C1 == 1) {
-    // For scalars, keep CTPOP if it is legal or custom.
-    if (!VT.isVector() && TLI.isOperationLegalOrCustom(ISD::CTPOP, CTVT))
+    // Keep the CTPOP if it is legal.
+    if (TLI.isOperationLegal(ISD::CTPOP, CTVT))
       return SDValue();
-    // This is based on X86's custom lowering for CTPOP which produces more
-    // instructions than the expansion here.
 
-    // (ctpop x) == 1 --> (x != 0) && ((x & x-1) == 0)
-    // (ctpop x) != 1 --> (x == 0) || ((x & x-1) != 0)
     SDValue Zero = DAG.getConstant(0, dl, CTVT);
     SDValue NegOne = DAG.getAllOnesConstant(dl, CTVT);
     assert(CTVT.isInteger());
     ISD::CondCode InvCond = ISD::getSetCCInverse(Cond, CTVT);
     SDValue Add = DAG.getNode(ISD::ADD, dl, CTVT, CTOp, NegOne);
     SDValue And = DAG.getNode(ISD::AND, dl, CTVT, CTOp, Add);
     SDValue LHS = DAG.getSetCC(dl, VT, CTOp, Zero, InvCond);
     SDValue RHS = DAG.getSetCC(dl, VT, And, Zero, Cond);

llvm/test/CodeGen/AArch64/arm64-popcnt.ll

 ; CHECK-NONEON-NEXT:    ret
   %cnt = tail call i64 @llvm.ctpop.i64(i64 %x)
   ret i64 %cnt
 }
 
 define i32 @ctpop_eq_one(i64 %x) nounwind readnone {
 ; CHECK-LABEL: ctpop_eq_one:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d0, x0
-; CHECK-NEXT:    cnt.8b v0, v0
-; CHECK-NEXT:    uaddlv.8b h0, v0
-; CHECK-NEXT:    fmov w8, s0
-; CHECK-NEXT:    cmp x8, #1
-; CHECK-NEXT:    cset w0, eq
+; CHECK-NEXT:    sub x8, x0, #1
+; CHECK-NEXT:    tst x0, x8
+; CHECK-NEXT:    ccmp x0, #0, #4, eq
+; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
 ;
 ; CHECK-NONEON-LABEL: ctpop_eq_one:
 ; CHECK-NONEON:       // %bb.0:
-; CHECK-NONEON-NEXT:    lsr x9, x0, #1
-; CHECK-NONEON-NEXT:    mov x8, #72340172838076673
-; CHECK-NONEON-NEXT:    and x9, x9, #0x5555555555555555
-; CHECK-NONEON-NEXT:    sub x9, x0, x9
-; CHECK-NONEON-NEXT:    lsr x10, x9, #2
-; CHECK-NONEON-NEXT:    and x9, x9, #0x3333333333333333
-; CHECK-NONEON-NEXT:    and x10, x10, #0x3333333333333333
-; CHECK-NONEON-NEXT:    add x9, x9, x10
-; CHECK-NONEON-NEXT:    add x9, x9, x9, lsr #4
-; CHECK-NONEON-NEXT:    and x9, x9, #0xf0f0f0f0f0f0f0f
-; CHECK-NONEON-NEXT:    mul x8, x9, x8
-; CHECK-NONEON-NEXT:    lsr x8, x8, #56
-; CHECK-NONEON-NEXT:    cmp x8, #1
-; CHECK-NONEON-NEXT:    cset w0, eq
+; CHECK-NONEON-NEXT:    sub x8, x0, #1
+; CHECK-NONEON-NEXT:    tst x0, x8
+; CHECK-NONEON-NEXT:    ccmp x0, #0, #4, eq
+; CHECK-NONEON-NEXT:    cset w0, ne
 ; CHECK-NONEON-NEXT:    ret
   %count = tail call i64 @llvm.ctpop.i64(i64 %x)
   %cmp = icmp eq i64 %count, 1
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 define i32 @ctpop_ne_one(i64 %x) nounwind readnone {
 ; CHECK-LABEL: ctpop_ne_one:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    fmov d0, x0
-; CHECK-NEXT:    cnt.8b v0, v0
-; CHECK-NEXT:    uaddlv.8b h0, v0
-; CHECK-NEXT:    fmov w8, s0
-; CHECK-NEXT:    cmp x8, #1
-; CHECK-NEXT:    cset w0, ne
+; CHECK-NEXT:    sub x8, x0, #1
+; CHECK-NEXT:    tst x0, x8
+; CHECK-NEXT:    ccmp x0, #0, #4, eq
+; CHECK-NEXT:    cset w0, eq
 ; CHECK-NEXT:    ret
 ;
 ; CHECK-NONEON-LABEL: ctpop_ne_one:
 ; CHECK-NONEON:       // %bb.0:
-; CHECK-NONEON-NEXT:    lsr x9, x0, #1
-; CHECK-NONEON-NEXT:    mov x8, #72340172838076673
-; CHECK-NONEON-NEXT:    and x9, x9, #0x5555555555555555
-; CHECK-NONEON-NEXT:    sub x9, x0, x9
-; CHECK-NONEON-NEXT:    lsr x10, x9, #2
-; CHECK-NONEON-NEXT:    and x9, x9, #0x3333333333333333
-; CHECK-NONEON-NEXT:    and x10, x10, #0x3333333333333333
-; CHECK-NONEON-NEXT:    add x9, x9, x10
-; CHECK-NONEON-NEXT:    add x9, x9, x9, lsr #4
-; CHECK-NONEON-NEXT:    and x9, x9, #0xf0f0f0f0f0f0f0f
-; CHECK-NONEON-NEXT:    mul x8, x9, x8
-; CHECK-NONEON-NEXT:    lsr x8, x8, #56
-; CHECK-NONEON-NEXT:    cmp x8, #1
-; CHECK-NONEON-NEXT:    cset w0, ne
+; CHECK-NONEON-NEXT:    sub x8, x0, #1
+; CHECK-NONEON-NEXT:    tst x0, x8
+; CHECK-NONEON-NEXT:    ccmp x0, #0, #4, eq
+; CHECK-NONEON-NEXT:    cset w0, eq
 ; CHECK-NONEON-NEXT:    ret
   %count = tail call i64 @llvm.ctpop.i64(i64 %x)
   %cmp = icmp ne i64 %count, 1
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 
 declare i32 @llvm.ctpop.i32(i32) nounwind readnone
 declare i64 @llvm.ctpop.i64(i64) nounwind readnone