[llvm] [X86] Keep scalar bf16/f16 selects in vector registers (PR #224218)

[Simon Pilgrim via llvm-commits]

================
@@ -48981,6 +48981,81 @@ static SDValue combineSelect(SDNode *N, SelectionDAG &DAG,
   bool CondConstantVector = ISD::isBuildVectorOfConstantSDNodes(Cond.getNode());
   unsigned EltBitWidth = VT.getScalarSizeInBits();
 
+  // Soft bf16/f16 scalar selects do a VSELECT in vector registers instead
+  // of a scalar CMOV, to avoid a GPR round-trip. Skip constant operands
+  // (cheaper as immediates) and compare-driven conds (CMOV already reuses
+  // the flags).
+  if (N->getOpcode() == ISD::SELECT && !CondVT.isVector() &&
+      Subtarget.hasSSE2() && !isIntOrFPConstant(LHS) &&
+      !isIntOrFPConstant(RHS)) {
+    // Only worth it if both operands already live in a vector register
+    auto IsBitcastFromGPR = [](SDValue Op) {
+      return Op.getOpcode() == ISD::BITCAST &&
+             Op.getOperand(0).getValueType().isScalarInteger();
+    };
+    SDValue F16LHS, F16RHS;
+    if (!VT.isVector() && isSoftF16(VT, Subtarget)) {
+      if (!IsBitcastFromGPR(LHS) || !IsBitcastFromGPR(RHS)) {
+        F16LHS = DAG.getBitcast(MVT::f16, LHS);
+        F16RHS = DAG.getBitcast(MVT::f16, RHS);
+      }
+    } else if (VT == MVT::i16 && LHS.getOpcode() == ISD::BITCAST &&
+               RHS.getOpcode() == ISD::BITCAST) {
+      MVT SVT = LHS.getOperand(0).getSimpleValueType();
+      if ((SVT == MVT::f16 || SVT == MVT::bf16) &&
+          SVT == RHS.getOperand(0).getSimpleValueType()) {
+        F16LHS = DAG.getBitcast(MVT::f16, LHS.getOperand(0));
+        F16RHS = DAG.getBitcast(MVT::f16, RHS.getOperand(0));
+      }
+    }
+
+    SDValue CondRoot = Cond;
+    while (CondRoot.getOpcode() == ISD::AND ||
----------------
RKSimon wrote:

Maybe tweak this into a IsFromSetCC helper?
```
auto IsFromSetCC = [](SDValue Cond) {
  while (Cond.getOpcode() == ISD::AND ||
         Cond.getOpcode() == ISD::ANY_EXTEND ||
         Cond.getOpcode() == ISD::ZERO_EXTEND ||
         Cond.getOpcode() == ISD::TRUNCATE)
      Cond = Cond.getOperand(0);
  return Cond.getOpcode() == ISD::SETCC || Cond.getOpcode() == X86ISD::SETCC;
}
if (F16LHS && !IsFromSetCC (Cond)) {
```
Maybe I'm just paranoid that CondRoot is going to get used by something....

https://github.com/llvm/llvm-project/pull/224218