[X86] Combine to vpdpbusd when operand is constant and small enough.

[LuoYuanke]

Update test cases.

[LuoYuanke]

Rebase.

[RKSimon]

Precommit the tests and rebase to show the diff in codegen?

llvm/test/CodeGen/X86/dpbusd_const.ll

2

Add avxvnni target coverage?

[LuoYuanke]

Address Simon's comments.

[RKSimon]

Pre-commit the tests + rebase so this patch shows the codegen diff?

[LuoYuanke]

Rebase and pre-checkin the test case.

[LuoYuanke]

llvm/test/CodeGen/X86/dpbusd_const.ll

124

This looks not good, as it accesses more memory.

[RKSimon]

llvm/test/CodeGen/X86/dpbusd_const.ll

124

This is very similar to https://llvm.org/PR45808 - we're not making use of implicit zeroing of upper subvectors to make smaller constants

[LuoYuanke]

Shrink the constant size to 128 bits if the higher bits are zeros.

[RKSimon]

LGTM - with one minor

llvm/lib/Target/X86/X86ISelLowering.cpp

9253

foled -> folded

[LuoYuanke]

Address Simon's comments. Thanks for review.

llvm/lib/Target/X86/X86ISelLowering.cpp

 ///  b. Repeated pattern of constants (e.g. <0,1,0,1> or <0,1,2,3,0,1,2,3>).
 ///
 /// The VBROADCAST node is returned when a pattern is found,
 /// or SDValue() otherwise.
 static SDValue lowerBuildVectorAsBroadcast(BuildVectorSDNode *BVOp,
                                            const X86Subtarget &Subtarget,
                                            SelectionDAG &DAG) {
   // VBROADCAST requires AVX.
   // TODO: Splats could be generated for non-AVX CPUs using SSE
   // instructions, but there's less potential gain for only 128-bit vectors.
   if (!Subtarget.hasAVX())
     return SDValue();
 
   MVT VT = BVOp->getSimpleValueType(0);
   unsigned NumElts = VT.getVectorNumElements();
   SDLoc dl(BVOp);
 

                          SDValue &Op1) {
   Op0 = Mul.getOperand(0);
   Op1 = Mul.getOperand(1);
 
   // The operand1 should be signed extend
   if (Op0.getOpcode() == ISD::SIGN_EXTEND)
     std::swap(Op0, Op1);
 
-  if (Op0.getOpcode() != ISD::ZERO_EXTEND)
-    return false;
-
   auto IsFreeTruncation = [](SDValue &Op) -> bool {
     if ((Op.getOpcode() == ISD::ZERO_EXTEND ||
          Op.getOpcode() == ISD::SIGN_EXTEND) &&
         Op.getOperand(0).getScalarValueSizeInBits() <= 8)
       return true;
 
-    // TODO: Support contant value.
-    return false;
+    auto *BV = dyn_cast<BuildVectorSDNode>(Op);
+    return (BV && BV->isConstant());
   };
 
   // (dpbusd (zext a), (sext, b)). Since the first operand should be unsigned
   // value, we need to check Op0 is zero extended value. Op1 should be signed
   // value, so we just check the signed bits.
   if ((IsFreeTruncation(Op0) &&
        DAG.computeKnownBits(Op0).countMaxActiveBits() <= 8) &&
       (IsFreeTruncation(Op1) && DAG.ComputeMaxSignificantBits(Op1) <= 8))

llvm/lib/Target/X86/X86PartialReduction.cpp

if (!ST->hasVNNI() && !ST->hasAVXVNNI())
     return false;
 
   Value *LHS = Mul->getOperand(0);
   Value *RHS = Mul->getOperand(1);
 
   if (isa<SExtInst>(LHS))
     std::swap(LHS, RHS);
 
-  if (!isa<ZExtInst>(LHS))
-    return false;
-
   auto IsFreeTruncation = [&](Value *Op) {
     if (auto *Cast = dyn_cast<CastInst>(Op)) {
       if (Cast->getParent() == Mul->getParent() &&
           (Cast->getOpcode() == Instruction::SExt ||
            Cast->getOpcode() == Instruction::ZExt) &&
           Cast->getOperand(0)->getType()->getScalarSizeInBits() <= 8)
         return true;
     }
-    // TODO: Support constant in ISel.
-    return false;
+
+    return isa<Constant>(Op);
   };
 
   // (dpbusd (zext a), (sext, b)). Since the first operand should be unsigned
   // value, we need to check LHS is zero extended value. RHS should be signed
   // value, so we just check the signed bits.
   if ((IsFreeTruncation(LHS) &&
        computeKnownBits(LHS, *DL).countMaxActiveBits() <= 8) &&
       (IsFreeTruncation(RHS) && ComputeMaxSignificantBits(RHS, *DL) <= 8))

llvm/test/CodeGen/X86/dpbusd_const.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512vnni -mattr=+avx512vl | FileCheck %s
+
+define i32 @mul_4xi8_zc_exceed(<4 x i8> %a, i32 %c) {
+; CHECK-LABEL: mul_4xi8_zc_exceed:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpmovzxbd {{.*#+}} xmm0 = xmm0[0],zero,zero,zero,xmm0[1],zero,zero,zero,xmm0[2],zero,zero,zero,xmm0[3],zero,zero,zero
+; CHECK-NEXT:    vpmaddwd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; CHECK-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vmovd %xmm0, %eax
+; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = zext <4 x i8> %a to <4 x i32>
+  %1 = mul nsw <4 x i32> %0, <i32 0, i32 1, i32 2, i32 128>
+  %2 = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %1)
+  %op.extra = add nsw i32 %2, %c
+  ret i32 %op.extra
+}
+
+define i32 @mul_4xi8_zc(<4 x i8> %a, i32 %c) {
+; CHECK-LABEL: mul_4xi8_zc:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; CHECK-NEXT:    vpblendw {{.*#+}} xmm0 = xmm0[0,1],xmm1[2,3,4,5,6,7]
+; CHECK-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; CHECK-NEXT:    vpdpbusd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm1
+; CHECK-NEXT:    vmovd %xmm1, %eax
+; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = zext <4 x i8> %a to <4 x i32>
+  %1 = mul nsw <4 x i32> %0, <i32 0, i32 1, i32 2, i32 127>
+  %2 = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %1)
+  %op.extra = add nsw i32 %2, %c
+  ret i32 %op.extra
+}
+
+define i32 @mul_4xi4_cz(<4 x i4> %a, i32 %c) {
+; CHECK-LABEL: mul_4xi4_cz:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpmovdb %xmm0, %xmm0
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vmovdqa {{.*#+}} xmm1 = [0,1,2,127,0,0,0,0,0,0,0,0,0,0,0,0]
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm0, %xmm1, %xmm2
+; CHECK-NEXT:    vmovd %xmm2, %eax
+; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = zext <4 x i4> %a to <4 x i32>
+  %1 = mul nsw <4 x i32> <i32 0, i32 1, i32 2, i32 127>, %0
+  %2 = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %1)
+  %op.extra = add nsw i32 %2, %c
+  ret i32 %op.extra
+}
+
+define i32 @mul_4xi8_cs(<4 x i8> %a, i32 %c) {
+; CHECK-LABEL: mul_4xi8_cs:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; CHECK-NEXT:    vpblendw {{.*#+}} xmm0 = xmm0[0,1],xmm1[2,3,4,5,6,7]
+; CHECK-NEXT:    vmovdqa {{.*#+}} xmm1 = [0,1,2,255,0,0,0,0,0,0,0,0,0,0,0,0]
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm0, %xmm1, %xmm2
+; CHECK-NEXT:    vmovd %xmm2, %eax
+; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = sext <4 x i8> %a to <4 x i32>
+  %1 = mul nsw <4 x i32> <i32 0, i32 1, i32 2, i32 255>, %0
+  %2 = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %1)
+  %op.extra = add nsw i32 %2, %c
+  ret i32 %op.extra
+}
+
+define i32 @mul_4xi8_cs_exceed(<4 x i8> %a, i32 %c) {
+; CHECK-LABEL: mul_4xi8_cs_exceed:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpmovsxbd %xmm0, %xmm0
+; CHECK-NEXT:    vpmaddwd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; CHECK-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vmovd %xmm0, %eax
+; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = sext <4 x i8> %a to <4 x i32>
+  %1 = mul nsw <4 x i32> <i32 0, i32 1, i32 2, i32 256>, %0
+  %2 = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %1)
+  %op.extra = add nsw i32 %2, %c
+  ret i32 %op.extra
+}
+
+define i32 @mul_16xi8_zc(<16 x i8> %a, i32 %c) {
+; CHECK-LABEL: mul_16xi8_zc:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; CHECK-NEXT:    vpdpbusd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm1
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; CHECK-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vmovd %xmm0, %eax
+; CHECK-NEXT:    addl %edi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = zext <16 x i8> %a to <16 x i32>
+  %1 = mul nsw <16 x i32> %0, <i32 0, i32 1, i32 2, i32 64, i32 0, i32 1, i32 2, i32 64, i32 0, i32 1, i32 2, i32 64, i32 0, i32 1, i32 2, i32 64>
+  %2 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %1)
+  %op.extra = add nsw i32 %2, %c
+  ret i32 %op.extra
+}
+
+declare i32 @llvm.vector.reduce.add.v4i32(<4 x i32>)
+declare i32 @llvm.vector.reduce.add.v16i32(<16 x i32>)