[llvm] [SDAG] Make Select-with-Identity-Fold More Flexible; NFC (PR #136554)

[Marius Kamp via llvm-commits]

https://github.com/mskamp created https://github.com/llvm/llvm-project/pull/136554

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.

>From 806fdc796828ba9940072c04dbd3621f0e56d0b5 Mon Sep 17 00:00:00 2001
From: Marius Kamp <msk at posteo.org>
Date: Fri, 18 Apr 2025 11:04:22 +0200
Subject: [PATCH] [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.
---
 llvm/include/llvm/CodeGen/TargetLowering.h    | 15 ++++++++++++++-
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 12 ++++++++----
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index 00c36266a069f..ace71bee0ac34 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/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
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index b175e35385ec6..c67614f4aa759 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/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);