[X86] Add support for calling SimplifyDemandedBits on the input of PDEP with a constant mask.

[craig.topper]

Summary

We can do several optimizations for PDEP using computeKnownBits and SimplifyDemandedBits

-If the MSBs of the output aren't demanded, those MSBs of the mask input aren't demanded either. We need to keep the most significant demanded bit of the mask and any mask bits before it.
-The number of possible ones in the mask determines how many bits of the lsbs of the other operand are demanded. Any bits of the mask we don't demand by the previous rule should not be counted.
-The result will have zeros in any position that the mask is zero.
-Since non-mask input bits can only be output in the original position or a higher bit position, the result will have at least as many trailing zeroes as the non-mask input.

Test cases have not been committed yet. But this patch was written to show the test diffs. Wanted to get feedback on the tests before committing.

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

49499

Can you move this functionality inside X86TargetLowering::SimplifyDemandedBitsForTargetNode?

Hopefully we can then expand it in the future to handle incoming DemandedBits masks as well to adjust the PDEP mask.

X86TargetLowering::computeKnownBitsForTargetNode support would be useful eventually as well.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

49499

Sure, but I don’t think it makes sense to move it without calculating the known bits based on the mask. So we’ll need more tests.

[craig.topper]

Check my logic

Demanded bits for mask = (1 << (bitwidth - lzcnt(input demanded bits)) - 1
Demanded bits for other operand = (1<< ctpop(~mask_known.zero & demanded_bits_for_mask)) - 1

Result known.zero = mask_known.zero

[craig.topper]

We can do several optimizations for PDEP using computeKnownBits and SimplifyDemandedBits

-If the MSBs of the output aren't demanded, those MSBs of the mask input aren't demanded either. We need to keep the most significant demanded bit of the mask and any mask bits before it.
-The number of possible ones in the mask determines how many bits of the lsbs of the other operand are demanded. Any bits of the mask we don't demand by the previous rule should not be counted.
-The result will have zeros in any position that the mask is zero.
-Since non-mask input bits can only be output in the original position or a higher bit position, the result will have at least as many trailing zeroes as the non-mask input.

Test cases have not been committed yet. But this patch was written to show the test diffs. Wanted to get feedback on the tests before committing.

[craig.topper]

Ping

[RKSimon]

LGTM - the logic makes sense to me - maybe add equivalent code into X86TTIImpl::simplifyDemandedUseBitsIntrinsic ? We already have pdep/pext constant folding support in instcombine

llvm/lib/Target/X86/X86ISelLowering.cpp

34064

It might be worth moving this into a KnownBits::DepositBits helper - IIRC RISCV has a similar instruction being finalised - maybe also move the instcombine constant folding handling into APInt ?

llvm/lib/Target/X86/X86ISelLowering.cpp

if (auto* Cst1 = dyn_cast<ConstantSDNode>(Op1)) {
       if ((Shift + Length) <= BitWidth) {
         Known = DAG.computeKnownBits(Op0, Depth + 1);
         Known = Known.extractBits(Length, Shift);
         Known = Known.zextOrTrunc(BitWidth);
       }
     }
     break;
   }
+  case X86ISD::PDEP: {
+    KnownBits Known2;
+    Known = DAG.computeKnownBits(Op.getOperand(1), DemandedElts, Depth + 1);
+    Known2 = DAG.computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
+    // Zeros are retained from the mask operand. But not ones.
+    Known.One.clearAllBits();
+    // The result will have at least as many trailing zeros as the non-mask
+    // operand since bits can only map to the same or higher bit position.
+    Known.Zero.setLowBits(Known2.countMinTrailingZeros());
+    break;
+  }
   case X86ISD::VTRUNC:
   case X86ISD::VTRUNCS:
   case X86ISD::VTRUNCUS:
   case X86ISD::CVTSI2P:
   case X86ISD::CVTUI2P:
   case X86ISD::CVTP2SI:
   case X86ISD::CVTP2UI:
   case X86ISD::MCVTP2SI:
case X86ISD::BEXTR: {
 
     // If the length is 0, replace with 0.
     KnownBits LengthBits = Known1.extractBits(8, 8);
     if (LengthBits.isZero())
       return TLO.CombineTo(Op, TLO.DAG.getConstant(0, SDLoc(Op), VT));
 
     break;
   }
+  case X86ISD::PDEP: {
+    SDValue Op0 = Op.getOperand(0);
+    SDValue Op1 = Op.getOperand(1);
+
+    unsigned DemandedBitsLZ = OriginalDemandedBits.countLeadingZeros();
+    APInt LoMask = APInt::getLowBitsSet(BitWidth, BitWidth - DemandedBitsLZ);
+
+    // If the demanded bits has leading zeroes, we don't demand those from the
+    // mask.
+    if (SimplifyDemandedBits(Op1, LoMask, Known, TLO, Depth + 1))
+      return true;
+
+    // The number of possible 1s in the mask determines the number of LSBs of
+    // operand 0 used. Undemanded bits from the mask don't matter so filter
+    // them before counting.
+    KnownBits Known2;
+    uint64_t Count = (~Known.Zero & LoMask).countPopulation();
+    APInt DemandedMask(APInt::getLowBitsSet(BitWidth, Count));
+    if (SimplifyDemandedBits(Op0, DemandedMask, Known2, TLO, Depth + 1))
+      return true;
+
+    // Zeroes are retained from the mask, but not ones.
+    Known.One.clearAllBits();
+    // The result will have at least as many trailing zeros as the non-mask
+    // operand since bits can only map to the same or higher bit position.
+    Known.Zero.setLowBits(Known2.countMinTrailingZeros());
+    return false;
+  }
   }
 
   return TargetLowering::SimplifyDemandedBitsForTargetNode(
       Op, OriginalDemandedBits, OriginalDemandedElts, Known, TLO, Depth);
 }
 
 SDValue X86TargetLowering::SimplifyMultipleUseDemandedBitsForTargetNode(
     SDValue Op, const APInt &DemandedBits, const APInt &DemandedElts,
if (NumElts < 8) {
     Src = DAG.getNode(ISD::CONCAT_VECTORS, dl, MVT::v8i16, Ops);
   }
 
   // Destination is vXf32 with at least 4 elements.
   EVT CvtVT = EVT::getVectorVT(*DAG.getContext(), MVT::f32,
                                std::max(4U, NumElts));
   SDValue Cvt, Chain;
   if (IsStrict) {
     Cvt = DAG.getNode(X86ISD::STRICT_CVTPH2PS, dl, {CvtVT, MVT::Other},
                       {N->getOperand(0), Src});
     Chain = Cvt.getValue(1);
   } else {
     Cvt = DAG.getNode(X86ISD::CVTPH2PS, dl, CvtVT, Src);
   }
 
   if (NumElts < 4) {
     assert(NumElts == 2 && "Unexpected size");
if (LN->isSimple()) {
       DAG.ReplaceAllUsesOfValueWith(SDValue(LN, 1), NewLd.getValue(1));
       return NewLd;
     }
   }
 
   return SDValue();
 }
 
+static SDValue combinePDEP(SDNode *N, SelectionDAG &DAG,
+                           TargetLowering::DAGCombinerInfo &DCI) {
+  unsigned NumBits = N->getSimpleValueType(0).getSizeInBits();
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  if (TLI.SimplifyDemandedBits(SDValue(N, 0),
+                               APInt::getAllOnesValue(NumBits), DCI))
+    return SDValue(N, 0);
+
+  return SDValue();
+}
+
 SDValue X86TargetLowering::PerformDAGCombine(SDNode *N,
                                              DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;
   switch (N->getOpcode()) {
   default: break;
   case ISD::SCALAR_TO_VECTOR:
     return combineScalarToVector(N, DAG);
   case ISD::EXTRACT_VECTOR_ELT:
SDValue X86TargetLowering::PerformDAGCombine(SDNode *N,
   case X86ISD::KSHIFTL:
   case X86ISD::KSHIFTR:     return combineKSHIFT(N, DAG, DCI);
   case ISD::FP16_TO_FP:     return combineFP16_TO_FP(N, DAG, Subtarget);
   case ISD::STRICT_FP_EXTEND:
   case ISD::FP_EXTEND:      return combineFP_EXTEND(N, DAG, Subtarget);
   case ISD::FP_ROUND:       return combineFP_ROUND(N, DAG, Subtarget);
   case X86ISD::VBROADCAST_LOAD: return combineVBROADCAST_LOAD(N, DAG, DCI);
   case X86ISD::MOVDQ2Q:     return combineMOVDQ2Q(N, DAG);
+  case X86ISD::PDEP:        return combinePDEP(N, DAG, DCI);
   }
 
   return SDValue();
 }
 
 bool X86TargetLowering::isTypeDesirableForOp(unsigned Opc, EVT VT) const {
   if (!isTypeLegal(VT))
     return false;

llvm/test/CodeGen/X86/bmi2-x86_64.ll

; CHECK-NEXT:    retq
   %y1 = load i64, i64* %y
   %tmp = tail call i64 @llvm.x86.bmi.pdep.64(i64 %x, i64 %y1)
   ret i64 %tmp
 }
 
 define i64 @pdep64_anyext(i32 %x)   {
 ; CHECK-LABEL: pdep64_anyext:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movslq %edi, %rax
-; CHECK-NEXT:    movabsq $6148914691236517205, %rcx # imm = 0x5555555555555555
-; CHECK-NEXT:    pdepq %rcx, %rax, %rax
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NEXT:    movabsq $6148914691236517205, %rax # imm = 0x5555555555555555
+; CHECK-NEXT:    pdepq %rax, %rdi, %rax
 ; CHECK-NEXT:    retq
   %x1 = sext i32 %x to i64
   %tmp = tail call i64 @llvm.x86.bmi.pdep.64(i64 %x1, i64 6148914691236517205)
   ret i64 %tmp
 }
 
 declare i64 @llvm.x86.bmi.pdep.64(i64, i64)
 

llvm/test/CodeGen/X86/bmi2.ll

 ; X86:       # %bb.0:
 ; X86-NEXT:    movswl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl $-1431655766, %ecx # imm = 0xAAAAAAAA
 ; X86-NEXT:    pdepl %ecx, %eax, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_anyext:
 ; X64:       # %bb.0:
-; X64-NEXT:    movswl %di, %eax
-; X64-NEXT:    movl $-1431655766, %ecx # imm = 0xAAAAAAAA
-; X64-NEXT:    pdepl %ecx, %eax, %eax
+; X64-NEXT:    movl $-1431655766, %eax # imm = 0xAAAAAAAA
+; X64-NEXT:    pdepl %eax, %edi, %eax
 ; X64-NEXT:    retq
   %x1 = sext i16 %x to i32
   %tmp = tail call i32 @llvm.x86.bmi.pdep.32(i32 %x1, i32 -1431655766)
   ret i32 %tmp
 }
 
 define i32 @pdep32_demandedbits(i32 %x) {
 ; X86-LABEL: pdep32_demandedbits:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl $1431655765, %ecx # imm = 0x55555555
 ; X86-NEXT:    pdepl %ecx, %eax, %eax
-; X86-NEXT:    andl $1431655765, %eax # imm = 0x55555555
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_demandedbits:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl $1431655765, %eax # imm = 0x55555555
 ; X64-NEXT:    pdepl %eax, %edi, %eax
-; X64-NEXT:    andl $1431655765, %eax # imm = 0x55555555
 ; X64-NEXT:    retq
   %tmp = tail call i32 @llvm.x86.bmi.pdep.32(i32 %x, i32 1431655765)
   %tmp2 = and i32 %tmp, 1431655765
   ret i32 %tmp2
 }
 
 define i32 @pdep32_demandedbits2(i32 %x, i32 %y) {
 ; X86-LABEL: pdep32_demandedbits2:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movzbl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    pdepl {{[0-9]+}}(%esp), %eax, %eax
 ; X86-NEXT:    andl $128, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_demandedbits2:
 ; X64:       # %bb.0:
-; X64-NEXT:    movzbl %dil, %eax
-; X64-NEXT:    pdepl %esi, %eax, %eax
+; X64-NEXT:    pdepl %esi, %edi, %eax
 ; X64-NEXT:    andl $128, %eax
 ; X64-NEXT:    retq
   %tmp = and i32 %x, 255
   %tmp2 = tail call i32 @llvm.x86.bmi.pdep.32(i32 %tmp, i32 %y)
   %tmp3 = and i32 %tmp2, 128
   ret i32 %tmp3
 }
 
 define i32 @pdep32_demandedbits_mask(i32 %x, i16 %y) {
 ; X86-LABEL: pdep32_demandedbits_mask:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movswl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    pdepl %eax, %ecx, %eax
 ; X86-NEXT:    andl $32768, %eax # imm = 0x8000
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_demandedbits_mask:
 ; X64:       # %bb.0:
-; X64-NEXT:    movswl %si, %eax
-; X64-NEXT:    pdepl %eax, %edi, %eax
+; X64-NEXT:    pdepl %esi, %edi, %eax
 ; X64-NEXT:    andl $32768, %eax # imm = 0x8000
 ; X64-NEXT:    retq
   %tmp = sext i16 %y to i32
   %tmp2 = tail call i32 @llvm.x86.bmi.pdep.32(i32 %x, i32 %tmp)
   %tmp3 = and i32 %tmp2, 32768
   ret i32 %tmp3
 }
 
 define i32 @pdep32_demandedbits_mask2(i32 %x, i16 %y) {
 ; X86-LABEL: pdep32_demandedbits_mask2:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movswl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %ecx
 ; X86-NEXT:    pdepl %eax, %ecx, %eax
 ; X86-NEXT:    movzwl %ax, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_demandedbits_mask2:
 ; X64:       # %bb.0:
-; X64-NEXT:    movswl %si, %eax
-; X64-NEXT:    pdepl %eax, %edi, %eax
+; X64-NEXT:    pdepl %esi, %edi, %eax
 ; X64-NEXT:    movzwl %ax, %eax
 ; X64-NEXT:    retq
   %tmp = sext i16 %y to i32
   %tmp2 = tail call i32 @llvm.x86.bmi.pdep.32(i32 %x, i32 %tmp)
   %tmp3 = and i32 %tmp2, 65535
   ret i32 %tmp3
 }
 
 define i32 @pdep32_knownbits(i32 %x) {
 ; X86-LABEL: pdep32_knownbits:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-NEXT:    movl $1431655765, %ecx # imm = 0x55555555
-; X86-NEXT:    pdepl %ecx, %eax, %ecx
-; X86-NEXT:    movl %ecx, %eax
-; X86-NEXT:    andl $1431655765, %eax # imm = 0x55555555
-; X86-NEXT:    imull %ecx, %eax
+; X86-NEXT:    pdepl %ecx, %eax, %eax
+; X86-NEXT:    imull %eax, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_knownbits:
 ; X64:       # %bb.0:
 ; X64-NEXT:    movl $1431655765, %eax # imm = 0x55555555
-; X64-NEXT:    pdepl %eax, %edi, %ecx
-; X64-NEXT:    movl %ecx, %eax
-; X64-NEXT:    andl $1431655765, %eax # imm = 0x55555555
-; X64-NEXT:    imull %ecx, %eax
+; X64-NEXT:    pdepl %eax, %edi, %eax
+; X64-NEXT:    imull %eax, %eax
 ; X64-NEXT:    retq
   %tmp = tail call i32 @llvm.x86.bmi.pdep.32(i32 %x, i32 1431655765)
   %tmp2 = and i32 %tmp, 1431655765
   %tmp3 = mul i32 %tmp, %tmp2
   ret i32 %tmp3
 }
 
 define i32 @pdep32_knownbits2(i32 %x, i32 %y) {
 ; X86-LABEL: pdep32_knownbits2:
 ; X86:       # %bb.0:
 ; X86-NEXT:    movl $-256, %eax
 ; X86-NEXT:    andl {{[0-9]+}}(%esp), %eax
-; X86-NEXT:    pdepl {{[0-9]+}}(%esp), %eax, %ecx
-; X86-NEXT:    movl %ecx, %eax
-; X86-NEXT:    andl $-256, %eax
-; X86-NEXT:    imull %ecx, %eax
+; X86-NEXT:    pdepl {{[0-9]+}}(%esp), %eax, %eax
+; X86-NEXT:    imull %eax, %eax
 ; X86-NEXT:    retl
 ;
 ; X64-LABEL: pdep32_knownbits2:
 ; X64:       # %bb.0:
 ; X64-NEXT:    andl $-256, %edi
-; X64-NEXT:    pdepl %esi, %edi, %ecx
-; X64-NEXT:    movl %ecx, %eax
-; X64-NEXT:    andl $-256, %eax
-; X64-NEXT:    imull %ecx, %eax
+; X64-NEXT:    pdepl %esi, %edi, %eax
+; X64-NEXT:    imull %eax, %eax
 ; X64-NEXT:    retq
   %tmp = and i32 %x, -256
   %tmp2 = tail call i32 @llvm.x86.bmi.pdep.32(i32 %tmp, i32 %y)
   %tmp3 = and i32 %tmp2, -256
   %tmp4 = mul i32 %tmp2, %tmp3
   ret i32 %tmp4
 }