[X86] Use vectorized i256 bit counts when we know the source originated from the vector unit (#171589)

Currently we only permit i256 CTTZ/CTLZ AVX512 lowering when the source
is loadable as GPR->FPU transition costs would outweigh the
vectorization benefit.

This patch checks for other cases where the source can avoid the GPR - a
mayFoldToVector helper checks for a bitcast originally from a vector
type, as well as constant values and the original mayFoldLoad check.

There will be other cases for the mayFoldToVector helper, but I've just
used this for CTTZ/CTLZ initially.

Author: RKSimon

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -2846,6 +2846,15 @@ bool X86::mayFoldIntoZeroExtend(SDValue Op) {
   return false;
 }
 
+// Return true if its cheap to bitcast this to a vector type.
+static bool mayFoldIntoVector(SDValue Op, const X86Subtarget &Subtarget) {
+  if (peekThroughBitcasts(Op).getValueType().isVector())
+    return true;
+  if (isa<ConstantSDNode>(Op) || isa<ConstantFPSDNode>(Op))
+    return true;
+  return X86::mayFoldLoad(Op, Subtarget);
+}
+
 static bool isLogicOp(unsigned Opcode) {
   // TODO: Add support for X86ISD::FAND/FOR/FXOR/FANDN with test coverage.
   return ISD::isBitwiseLogicOp(Opcode) || X86ISD::ANDNP == Opcode;
@@ -33958,7 +33967,7 @@ void X86TargetLowering::ReplaceNodeResults(SDNode *N,
     EVT VT = N->getValueType(0);
     assert(Subtarget.hasCDI() && "AVX512CD required");
     assert((VT == MVT::i256 || VT == MVT::i512) && "Unexpected VT!");
-    if (VT == MVT::i256 && !X86::mayFoldLoad(N0, Subtarget))
+    if (VT == MVT::i256 && !mayFoldIntoVector(N0, Subtarget))
       return;
 
     unsigned SizeInBits = VT.getSizeInBits();