[SelectionDAG] Add DoNotPoisonEltMask to SimplifyMultipleUseDemandedBitsForTargetNode

Author: bjope

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -4388,6 +4388,7 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
   /// bitwise ops etc.
   virtual SDValue SimplifyMultipleUseDemandedBitsForTargetNode(
       SDValue Op, const APInt &DemandedBits, const APInt &DemandedElts,
+      const APInt &DoNotPoisonEltMask,
       SelectionDAG &DAG, unsigned Depth) const;
 
   /// Return true if this function can prove that \p Op is never poison

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -1000,7 +1000,7 @@ SDValue TargetLowering::SimplifyMultipleUseDemandedBits(
 
     if (Op.getOpcode() >= ISD::BUILTIN_OP_END)
       if (SDValue V = SimplifyMultipleUseDemandedBitsForTargetNode(
-              Op, DemandedBits, DemandedElts | DoNotPoisonEltMask, DAG, Depth))
+              Op, DemandedBits, DemandedElts, DoNotPoisonEltMask, DAG, Depth))
         return V;
     break;
   }
@@ -4058,12 +4058,14 @@ bool TargetLowering::SimplifyDemandedBitsForTargetNode(
           Op.getOpcode() == ISD::INTRINSIC_VOID) &&
          "Should use SimplifyDemandedBits if you don't know whether Op"
          " is a target node!");
-  computeKnownBitsForTargetNode(Op, Known, DemandedElts, TLO.DAG, Depth);
+  computeKnownBitsForTargetNode(Op, Known, DemandedElts,
+                                TLO.DAG, Depth);
   return false;
 }
 
 SDValue TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
     SDValue Op, const APInt &DemandedBits, const APInt &DemandedElts,
+    const APInt &DoNotPoisonEltMask,
     SelectionDAG &DAG, unsigned Depth) const {
   assert(
       (Op.getOpcode() >= ISD::BUILTIN_OP_END ||

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -43672,9 +43672,9 @@ bool X86TargetLowering::SimplifyDemandedVectorEltsForTargetNode(
       return true;
 
     if (!DemandedElts.isAllOnes()) {
-      SDValue NewLHS = SimplifyMultipleUseDemandedBits(LHS, BitsLHS, EltsLHS | DoNotPoisonEltMask,
+      SDValue NewLHS = SimplifyMultipleUseDemandedBits(LHS, BitsLHS, EltsLHS, DoNotPoisonEltMask,
                                                        TLO.DAG, Depth + 1);
-      SDValue NewRHS = SimplifyMultipleUseDemandedBits(RHS, BitsRHS, EltsRHS | DoNotPoisonEltMask,
+      SDValue NewRHS = SimplifyMultipleUseDemandedBits(RHS, BitsRHS, EltsRHS, DoNotPoisonEltMask,
                                                        TLO.DAG, Depth + 1);
       if (NewLHS || NewRHS) {
         NewLHS = NewLHS ? NewLHS : LHS;
@@ -43796,6 +43796,7 @@ bool X86TargetLowering::SimplifyDemandedVectorEltsForTargetNode(
   case X86ISD::BLENDI: {
     SmallVector<int, 16> BlendMask;
     DecodeBLENDMask(NumElts, Op.getConstantOperandVal(2), BlendMask);
+    // TODO: Do we really need to consider the DoNotPoisonEltMask here?
     if (SDValue R = combineBlendOfPermutes(
             VT.getSimpleVT(), Op.getOperand(0), Op.getOperand(1), BlendMask,
             DemandedElts | DoNotPoisonEltMask, TLO.DAG, Subtarget, SDLoc(Op)))
@@ -44906,6 +44907,7 @@ bool X86TargetLowering::SimplifyDemandedBitsForTargetNode(
 
 SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
     SDValue Op, const APInt &DemandedBits, const APInt &DemandedElts,
+    const APInt &DoNotPoisonEltMask,
     SelectionDAG &DAG, unsigned Depth) const {
   int NumElts = DemandedElts.getBitWidth();
   unsigned Opc = Op.getOpcode();
@@ -44919,7 +44921,8 @@ SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
     auto *CIdx = dyn_cast<ConstantSDNode>(Op.getOperand(2));
     MVT VecVT = Vec.getSimpleValueType();
     if (CIdx && CIdx->getAPIntValue().ult(VecVT.getVectorNumElements()) &&
-        !DemandedElts[CIdx->getZExtValue()])
+        !DemandedElts[CIdx->getZExtValue()] &&
+        !DoNotPoisonEltMask[CIdx->getZExtValue()])
       return Vec;
     break;
   }
@@ -44954,7 +44957,7 @@ SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
     SDValue LHS = Op.getOperand(1);
     SDValue RHS = Op.getOperand(2);
 
-    KnownBits CondKnown = DAG.computeKnownBits(Cond, DemandedElts, Depth + 1);
+    KnownBits CondKnown = DAG.computeKnownBits(Cond, DemandedElts | DoNotPoisonEltMask, Depth + 1);
     if (CondKnown.isNegative())
       return LHS;
     if (CondKnown.isNonNegative())
@@ -44966,8 +44969,8 @@ SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
     SDValue LHS = Op.getOperand(0);
     SDValue RHS = Op.getOperand(1);
 
-    KnownBits LHSKnown = DAG.computeKnownBits(LHS, DemandedElts, Depth + 1);
-    KnownBits RHSKnown = DAG.computeKnownBits(RHS, DemandedElts, Depth + 1);
+    KnownBits LHSKnown = DAG.computeKnownBits(LHS, DemandedElts | DoNotPoisonEltMask, Depth + 1);
+    KnownBits RHSKnown = DAG.computeKnownBits(RHS, DemandedElts | DoNotPoisonEltMask, Depth + 1);
 
     // If all of the demanded bits are known 0 on LHS and known 0 on RHS, then
     // the (inverted) LHS bits cannot contribute to the result of the 'andn' in
@@ -44981,7 +44984,8 @@ SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
   APInt ShuffleUndef, ShuffleZero;
   SmallVector<int, 16> ShuffleMask;
   SmallVector<SDValue, 2> ShuffleOps;
-  if (getTargetShuffleInputs(Op, DemandedElts, ShuffleOps, ShuffleMask,
+  if (getTargetShuffleInputs(Op, DemandedElts | DoNotPoisonEltMask,
+                             ShuffleOps, ShuffleMask,
                              ShuffleUndef, ShuffleZero, DAG, Depth, false)) {
     // If all the demanded elts are from one operand and are inline,
     // then we can use the operand directly.
@@ -45000,7 +45004,7 @@ SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
       APInt IdentityOp = APInt::getAllOnes(NumOps);
       for (int i = 0; i != NumElts; ++i) {
         int M = ShuffleMask[i];
-        if (!DemandedElts[i] || ShuffleUndef[i])
+        if (!(DemandedElts[i] || DoNotPoisonEltMask[i]) || ShuffleUndef[i])
           continue;
         int OpIdx = M / NumElts;
         int EltIdx = M % NumElts;
@@ -45021,7 +45025,7 @@ SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
   }
 
   return TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
-      Op, DemandedBits, DemandedElts, DAG, Depth);
+      Op, DemandedBits, DemandedElts, DoNotPoisonEltMask, DAG, Depth);
 }
 
 bool X86TargetLowering::isGuaranteedNotToBeUndefOrPoisonForTargetNode(

llvm/lib/Target/X86/X86ISelLowering.h

@@ -1325,6 +1325,7 @@ namespace llvm {
 
     SDValue SimplifyMultipleUseDemandedBitsForTargetNode(
         SDValue Op, const APInt &DemandedBits, const APInt &DemandedElts,
+        const APInt &DoNotPoisonEltMask,
         SelectionDAG &DAG, unsigned Depth) const override;
 
     bool isGuaranteedNotToBeUndefOrPoisonForTargetNode(