[X86] Distribute Certain Bitwise Operations over SELECT

InstCombine canonicalizes `(select P (and X (- X)) X)` to
`(and (select P (- X) umax) X)`. This is counterproductive for the X86
backend when BMI is available because we can encode `(and X (- X))`
using the `BLSI` instruction. A similar situation arises if we have
`(select P (and X (sub X 1)) X)` (prevents use of `BLSR` instruction) or
`(select P (xor X (sub X 1)) X)` (prevents use of `BLSMSK` instruction).

Trigger the inverse transformation in the X86 backend if BMI is
available and we can use the mentioned BMI instructions. This is done by
overriding the appropriate `shouldFoldSelectWithIdentityConstant()`
overload. In this way, we get `(select P (and X (- X)) X)` again, which
enables the use of `BLSI` (similar for the other cases described above).

Alive proofs: https://alive2.llvm.org/ce/z/MT_pKi

Fixes #131587, fixes #133848.

Author: mskamp

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -27,6 +27,7 @@
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/VectorUtils.h"
+#include "llvm/CodeGen/ISDOpcodes.h"
 #include "llvm/CodeGen/IntrinsicLowering.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
@@ -35385,18 +35386,48 @@ bool X86TargetLowering::isNarrowingProfitable(SDNode *N, EVT SrcVT,
 
 bool X86TargetLowering::shouldFoldSelectWithIdentityConstant(unsigned Opcode,
                                                              EVT VT) const {
+  if (!VT.isVector()) {
+    if (!Subtarget.hasBMI() || (VT != MVT::i32 && VT != MVT::i64))
+      return false;
+    if (Opcode != ISD::AND && Opcode != ISD::XOR)
+      return false;
+    return true;
+  }
+
   // TODO: This is too general. There are cases where pre-AVX512 codegen would
   //       benefit. The transform may also be profitable for scalar code.
   if (!Subtarget.hasAVX512())
     return false;
   if (!Subtarget.hasVLX() && !VT.is512BitVector())
     return false;
-  if (!VT.isVector() || VT.getScalarType() == MVT::i1)
+  if (VT.getScalarType() == MVT::i1)
     return false;
 
   return true;
 }
 
+bool X86TargetLowering::shouldFoldSelectWithIdentityConstant(
+    unsigned BinOpcode, unsigned SelectOpcode, SDValue X,
+    SDValue NonIdConstNode) const {
+  using namespace llvm::SDPatternMatch;
+
+  if (SelectOpcode == ISD::VSELECT)
+    return true;
+  // BLSI
+  if (BinOpcode == ISD::AND && sd_match(NonIdConstNode, m_Neg(m_Specific(X))))
+    return true;
+  // BLSR
+  if (BinOpcode == ISD::AND &&
+      sd_match(NonIdConstNode, m_Add(m_Specific(X), m_AllOnes())))
+    return true;
+  // BLSMSK
+  if (BinOpcode == ISD::XOR &&
+      sd_match(NonIdConstNode, m_Add(m_Specific(X), m_AllOnes())))
+    return true;
+
+  return false;
+}
+
 /// Targets can use this to indicate that they only support *some*
 /// VECTOR_SHUFFLE operations, those with specific masks.
 /// By default, if a target supports the VECTOR_SHUFFLE node, all mask values

llvm/lib/Target/X86/X86ISelLowering.h

@@ -1463,6 +1463,11 @@ namespace llvm {
     bool shouldFoldSelectWithIdentityConstant(unsigned BinOpcode,
                                               EVT VT) const override;
 
+    bool
+    shouldFoldSelectWithIdentityConstant(unsigned BinOpcode,
+                                         unsigned SelectOpcode, SDValue X,
+                                         SDValue NonIdConstNode) const override;
+
     /// Given an intrinsic, checks if on the target the intrinsic will need to map
     /// to a MemIntrinsicNode (touches memory). If this is the case, it returns
     /// true and stores the intrinsic information into the IntrinsicInfo that was