[SDAG] Make Select-with-Identity-Fold More Flexible; NFC

This change introduces a new overload of the method
`shouldFoldSelectWithIdentityConstant()`, which takes the opcode of the
select node and the non-identity operand of the select node and is
called after all other checks have been performed. Moreover, this change
adjusts the precondition of the fold so that it would work for `SELECT`
nodes in addition to `VSELECT` nodes.

No functional change is intended because the default (and currently
only) implementation restricts the fold to a `VSELECT` node; the same
restriction as before.

The rationale of this change is to make more fine grained decisions
possible when to revert the InstCombine canonicalization of
`(select c (binop x y) y)` to `(binop (select c x idc) y)` in the
backends.

Author: mskamp

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -3351,13 +3351,26 @@ class TargetLoweringBase {
   }
 
   /// Return true if pulling a binary operation into a select with an identity
-  /// constant is profitable. This is the inverse of an IR transform.
+  /// constant is profitable for the given binary operation and type. This is
+  /// the inverse of an IR transform.
   /// Example: X + (Cond ? Y : 0) --> Cond ? (X + Y) : X
   virtual bool shouldFoldSelectWithIdentityConstant(unsigned BinOpcode,
                                                     EVT VT) const {
     return false;
   }
 
+  /// Return true if pulling a binary operation into a select with an identity
+  /// constant is profitable for the given binary operation, select operation,
+  /// operand to the binary operation, and select operand that is not the
+  /// identity constant. This is a more fine-grained variant of the previous
+  /// overload that is called only if the previous overload returned true.
+  virtual bool
+  shouldFoldSelectWithIdentityConstant(unsigned BinOpcode,
+                                       unsigned SelectOpcode, SDValue X,
+                                       SDValue NonIdConstNode) const {
+    return SelectOpcode == ISD::VSELECT;
+  }
+
   /// Return true if it is beneficial to convert a load of a constant to
   /// just the constant itself.
   /// On some targets it might be more efficient to use a combination of

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -2425,8 +2425,9 @@ static SDValue foldSelectWithIdentityConstant(SDNode *N, SelectionDAG &DAG,
   if (ShouldCommuteOperands)
     std::swap(N0, N1);
 
-  // TODO: Should this apply to scalar select too?
-  if (N1.getOpcode() != ISD::VSELECT || !N1.hasOneUse())
+  unsigned SelOpcode = N1.getOpcode();
+  if ((SelOpcode != ISD::VSELECT && SelOpcode != ISD::SELECT) ||
+      !N1.hasOneUse())
     return SDValue();
 
   // We can't hoist all instructions because of immediate UB (not speculatable).
@@ -2439,17 +2440,20 @@ static SDValue foldSelectWithIdentityConstant(SDNode *N, SelectionDAG &DAG,
   SDValue Cond = N1.getOperand(0);
   SDValue TVal = N1.getOperand(1);
   SDValue FVal = N1.getOperand(2);
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
   // This transform increases uses of N0, so freeze it to be safe.
   // binop N0, (vselect Cond, IDC, FVal) --> vselect Cond, N0, (binop N0, FVal)
   unsigned OpNo = ShouldCommuteOperands ? 0 : 1;
-  if (isNeutralConstant(Opcode, N->getFlags(), TVal, OpNo)) {
+  if (isNeutralConstant(Opcode, N->getFlags(), TVal, OpNo) &&
+      TLI.shouldFoldSelectWithIdentityConstant(Opcode, SelOpcode, N0, FVal)) {
     SDValue F0 = DAG.getFreeze(N0);
     SDValue NewBO = DAG.getNode(Opcode, SDLoc(N), VT, F0, FVal, N->getFlags());
     return DAG.getSelect(SDLoc(N), VT, Cond, F0, NewBO);
   }
   // binop N0, (vselect Cond, TVal, IDC) --> vselect Cond, (binop N0, TVal), N0
-  if (isNeutralConstant(Opcode, N->getFlags(), FVal, OpNo)) {
+  if (isNeutralConstant(Opcode, N->getFlags(), FVal, OpNo) &&
+      TLI.shouldFoldSelectWithIdentityConstant(Opcode, SelOpcode, N0, TVal)) {
     SDValue F0 = DAG.getFreeze(N0);
     SDValue NewBO = DAG.getNode(Opcode, SDLoc(N), VT, F0, TVal, N->getFlags());
     return DAG.getSelect(SDLoc(N), VT, Cond, NewBO, F0);