[SelectionDAG][X86] Enable SimplifySetCC CTPOP transforms for vector splats

This enables these transforms for vectors:
(ctpop x) u< 2 -> (x & x-1) == 0
(ctpop x) u> 1 -> (x & x-1) != 0
(ctpop x) == 1 --> (x != 0) && ((x & x-1) == 0)
(ctpop x) != 1 --> (x == 0) || ((x & x-1) != 0)

All enabled if CTPOP isn't Legal. This differs from the scalar
behavior where the first two are done unconditionally and the
last two are done if CTPOP isn't Legal or Custom. The Legal
check produced better results for vectors based on X86's
custom handling. Might be worth re-visiting scalars here.

I disabled the looking through truncate for vectors. The
code that creates new setcc can use the same result VT as the
original setcc even if we truncated the input. That may work
work for most scalars, but definitely wouldn't work for vectors
unless it was a vector of i1.

Fixes or at least improves PR47825

Reviewed By: spatel

Differential Revision: https://reviews.llvm.org/D89346

Author: topperc

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -3383,12 +3383,11 @@ static SDValue simplifySetCCWithCTPOP(const TargetLowering &TLI, EVT VT,
                                       SDValue N0, const APInt &C1,
                                       ISD::CondCode Cond, const SDLoc &dl,
                                       SelectionDAG &DAG) {
-  assert(!VT.isVector() && "Vectors are not supported yet!");
-
   // Look through truncs that don't change the value of a ctpop.
+  // FIXME: Add vector support? Need to be careful with setcc result type below.
   SDValue CTPOP = N0;
-  if (N0.getOpcode() == ISD::TRUNCATE && N0.hasOneUse() &&
-      N0.getValueSizeInBits() > Log2_32(N0.getOperand(0).getValueSizeInBits()))
+  if (N0.getOpcode() == ISD::TRUNCATE && N0.hasOneUse() && !VT.isVector() &&
+      N0.getScalarValueSizeInBits() > Log2_32(N0.getOperand(0).getScalarValueSizeInBits()))
     CTPOP = N0.getOperand(0);
 
   if (CTPOP.getOpcode() != ISD::CTPOP || !CTPOP.hasOneUse())
@@ -3400,6 +3399,13 @@ static SDValue simplifySetCCWithCTPOP(const TargetLowering &TLI, EVT VT,
   // (ctpop x) u< 2 -> (x & x-1) == 0
   // (ctpop x) u> 1 -> (x & x-1) != 0
   if ((Cond == ISD::SETULT && C1 == 2) || (Cond == ISD::SETUGT && C1 == 1)) {
+    // If this is a vector CTPOP, keep the CTPOP if it is legal.
+    // This based on X86's custom lowering for vector CTPOP which produces more
+    // instructions than the expansion here.
+    // TODO: Should we check if CTPOP is legal(or custom) for scalars?
+    if (VT.isVector() && TLI.isOperationLegal(ISD::CTPOP, CTVT))
+      return SDValue();
+
     SDValue NegOne = DAG.getAllOnesConstant(dl, CTVT);
     SDValue Add = DAG.getNode(ISD::ADD, dl, CTVT, CTOp, NegOne);
     SDValue And = DAG.getNode(ISD::AND, dl, CTVT, CTOp, Add);
@@ -3408,8 +3414,16 @@ static SDValue simplifySetCCWithCTPOP(const TargetLowering &TLI, EVT VT,
   }
 
   // If ctpop is not supported, expand a power-of-2 comparison based on it.
-  if ((Cond == ISD::SETEQ || Cond == ISD::SETNE) && C1 == 1 &&
-      !TLI.isOperationLegalOrCustom(ISD::CTPOP, CTVT)) {
+  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))
+      return SDValue();
+    // For vectors, keep CTPOP only if it is legal.
+    // This is based on X86's custom lowering for CTPOP which produces more
+    // instructions than the expansion here.
+    if (VT.isVector() && TLI.isOperationLegal(ISD::CTPOP, CTVT))
+      return SDValue();
+
     // (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);
@@ -3460,6 +3474,15 @@ SDValue TargetLowering::SimplifySetCC(EVT VT, SDValue N0, SDValue N1,
       !DAG.getNodeIfExists(ISD::SUB, DAG.getVTList(OpVT), { N0, N1 } ))
     return DAG.getSetCC(dl, VT, N1, N0, SwappedCC);
 
+  if (auto *N1C = isConstOrConstSplat(N1)) {
+    const APInt &C1 = N1C->getAPIntValue();
+
+    // Optimize some CTPOP cases.
+    if (SDValue V = simplifySetCCWithCTPOP(*this, VT, N0, C1, Cond, dl, DAG))
+      return V;
+  }
+
+  // FIXME: Support vectors.
   if (auto *N1C = dyn_cast<ConstantSDNode>(N1.getNode())) {
     const APInt &C1 = N1C->getAPIntValue();
 
@@ -3487,10 +3510,6 @@ SDValue TargetLowering::SimplifySetCC(EVT VT, SDValue N0, SDValue N1,
       }
     }
 
-    // Optimize some CTPOP cases.
-    if (SDValue V = simplifySetCCWithCTPOP(*this, VT, N0, C1, Cond, dl, DAG))
-      return V;
-
     // (zext x) == C --> x == (trunc C)
     // (sext x) == C --> x == (trunc C)
     if ((Cond == ISD::SETEQ || Cond == ISD::SETNE) &&