[X86] Combine reduce (add (mul x, y)) to VNNI instruction.

[LuoYuanke]

Summary

For below C code, we can use VNNI to combine the mul and add operation.

int usdot_prod_qi(unsigned char *restrict a, char *restrict b, int c, int n) {

int i;
for (i = 0; i < 32; i++) {
  c += ((int)a[i] * (int)b[i]);
}
return c;

}

We didn't support the combine acoss basic block in this patch.

[lebedev.ri]

Any explicit checks for extension/truncation and their bitwidth delta instantly make me suspicious nowadays.
Does this deal with commutativity?
I think what you want to check is the number of known sign bits / known leading zero bits.

[LuoYuanke]

update the title.

[LuoYuanke]

In D116039#3202764, @lebedev.ri wrote:

Any explicit checks for extension/truncation and their bitwidth delta instantly make me suspicious nowadays.

Pls comments on the code, so that I can easily understand.

Does this deal with commutativity?

At X86ISelLowering.cpp : 41763

I think what you want to check is the number of known sign bits / known leading zero bits.

Pls comments on the code, so that I can easily understand.

[lebedev.ri]

llvm/lib/Target/X86/X86ISelLowering.cpp

41756–41774

Any explicit checks for extension/truncation and their bitwidth delta instantly make me suspicious nowadays.
Does this deal with commutativity?
I think what you want to check is the number of known sign bits / known leading zero bits.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

41756–41774

If the sign extend side can be proven to be positive, the sign extend might be hidden as zero extend. This is why tryMAddReplacement checks for FreeTruncations and calls ComputeNumSignBits.

41817

Opcode is only used once. Why not just make it part of the getNode call?

[LuoYuanke]

Address Craig and Roman's comments.

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

41756–41774

Thanks Craig and Roman. I enhanced it in the new patch.

[LuoYuanke]

Remove debug code.

[lebedev.ri]

Please fix the patch description - both extensions there are signed, is that actually the specification for the intrinsic?

llvm/lib/Target/X86/X86ISelLowering.cpp

41790

I'm not sure i follow.
Why is this okay with negative numbers?

41790–41791

This still does not handle the commutative variant.

41791

You want ComputeMinSignedBits() <= 8 to check for sext-like

[lebedev.ri]

llvm/lib/Target/X86/X86ISelLowering.cpp

41790

https://alive2.llvm.org/ce/z/WGxXrz vs https://alive2.llvm.org/ce/z/UbWhNv

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

41812

createVPDPBUSD ?

42108

Should this be called combineVPDPBUSDPattern? VNNI is the ISA no?

llvm/test/CodeGen/X86/dpbusd.ll

4

You might be able to use a common prefix for some of these to reduce check duplication

6

Drop dso_local?

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

41790

I'm not sure i follow.
Why is this okay with negative numbers?

Here is description for VPDPBUSD from https://software.intel.com/content/www/us/en/develop/download/intel-architecture-instruction-set-extensions-programming-reference.html?wapkw=instruction. The first operand is unsigned, and the second operand is signed.

Multiplies the individual unsigned bytes of the first source operand by the corresponding signed bytes of the second source operand, producing intermediate signed word results. The word results are then summed and accumulated in the destination dword element size operand.

41790–41791

This still does not handle the commutative variant.

Sorry, I don't understand it very well. I do it with " std::swap(Op0, Op1)" in line 41764;

41791

You want ComputeMinSignedBits() <= 8 to check for sext-like

[LuoYuanke]

Address Roman and Simon's comments. Thanks for the review.

[lebedev.ri]

Could you please explain why there are both the knownbits-based checks, and checks for ISD::SIGN/ZERO_EXTEND nodes?

[LuoYuanke]

In D116039#3206040, @lebedev.ri wrote:

Could you please explain why there are both the knownbits-based checks, and checks for ISD::SIGN/ZERO_EXTEND nodes?

The VPDPBUSD multiplies the individual unsigned bytes of the first source operand by the corresponding signed bytes of the second source operand, producing intermediate signed word results. The word results are then summed and accumulated in the destination dword element size operand.

For src2, it is signed value, so we don't need to check for ISD::SIGN nodes, because if the signed bits are 1 it is negative value and if the signed bits are 0 it is positive value.
But for src1, it is unsigned value. If it is a positive value it is OK, but if it is a negative value we can't use VPDPBUSD to combine the original nodes. See test case mul_sext_i4i4 and mul_zext_i4i4() in dpbusd_i4.ll. For mul_zext_i4i4 we can use VPDPBUSD, but for mul_sext_i4i4 we can't because the src1 may be negative value.

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

41780

Maybe we can remove IsFreeTruncation() check as Roman mentions.
Roman, do you mean to remove IsFreeTruncation() check?

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

41780

Maybe we can remove IsFreeTruncation() check as Roman mentions.
Roman, do you mean to remove IsFreeTruncation() check?

If I remove the ISD::SIGN/ZERO_EXTEND check, I got crash with below test case in createVPDPBUSD(). I think there is room to improve the patch to cover more pattern. But to be conservatively I'd like to improve it in another patch, so that if we have regression we can revert less code.
Hi Roman,
What do you think?

declare i32 @llvm.vector.reduce.add.v16i32(<16 x i32>)

define dso_local i32 @mul_i4i2(<16 x i4> %b, i32 %c) {
entry:
  %0 = trunc <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15> to <16 x i16>
  %1 = zext <16 x i16> %0 to <16 x i32>
  %2 = zext <16 x i4> %b to <16 x i32>
  %3 = mul nsw <16 x i32> %2, %1
  %4 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %3)
  %op.extra = add nsw i32 %4, %c
  ret i32 %op.extra
}

[LuoYuanke]

Rebase.

[pengfei]

Format the code as Lint suggested?

llvm/lib/Target/X86/X86ISelLowering.cpp

41810

Can we return false here if Op0 is not ZERO_EXTEND?

41813–41814

How about ANY_EXTEND ? The same below.

41826

It's better to hoist to line 41810. How about ANY_EXTEND?

41870

It is always 2? Better to add a comment to explain.

41887–41888

AVXVNNI implies AVX2. We won't need to split to 128 bits.

42156–42157

Can we generate i32 first then do the truncation?

42162–42163

Can check DCI.isAfterLegalizeDAG() before calling the function instead?

42210

Can ues ExtractVT.getSizeInBits() derectly.

llvm/test/CodeGen/X86/dpbusd.ll

14

This is the only and a strange diff with the AVX512 code. Is there anything wrong in one of each?

[craig.topper]

Are we missing test cases for 32 x i8 and 64 x i8?

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

42162–42163

I don't think that works. We should be able to handle 32 x i8 and 64 x i8 which would have zero_extend and sign_extend with illegal result types.

[LuoYuanke]

I'll update the patch according to Phoebe and Craig's comments.

llvm/lib/Target/X86/X86ISelLowering.cpp

41813–41814

This check the opcode, so we need check both zero extend and sign extend. I'm not sure if any extend also works, because the upper bits is undefined. What's the signed bit for any extend?

42156–42157

I'm not sure if the result overflow, truncating back to i16 or some other types remain the same value. How about leave it as an enhancement?

llvm/test/CodeGen/X86/dpbusd.ll

14

This test doesn't generate vpdpbusd instruction, so the AVX512VNNI and AVX512VL generate the same code. For other test case, AVX512VNNI can only use zmm register, but AVX512VNNI + AVX512VL can use xmm register.

[LuoYuanke]

Address Phoebe and Craig's comments.

[LuoYuanke]

llvm/test/CodeGen/X86/dpbusd.ll

14

For vpdpbusd_64xi32, the result is the same.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

41813–41814

It is undefined and ComputeMinSignedBits will return BitWidth - 1 for it.

41818

You can use Op.getOperand(0).getScalarValueSizeInBits() to simplify this

41828

Can we use DAG.computeKnownBits(Op0).countMaxActiveBits() <= 8 to make this more readable?

41856

Why are Ext0 and Ext1 passed by const reference? SDValue should be passed by value.

41864

Can we just TRUNCATE the Ext nodes without assuming they are extend nodes. That way it just works when you support constants in the future?

41870

Can we do this without assuming the node is a SIGN/ZERO_EXTEND? Just truncate the original node to Vi8VT.

42164

This code isn't handling vpdpwssd so why mention it here?

42176

Is this code valid for this transform? There's a large comment of justification for why it is ok for SAD. I think I only saw a test for the SIGN_EXTEND case?

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

42176

Oops I see the other test. I need to think about the math.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

42176

I don't think we can do this if the multiply result is zero extended. Each of the 4 multiplies done by vpdpbusd compute a signed 16-bit product that will be sign extended before adding into the accumulator.

I think we also need to verify that the multiply has at least 2x the number of bits of the input. We shouldn't match (sign_extend (mul (vXi9 (zext (vXi8 X))), (vXi9 (zext (vXi8 Y)))). Does anything prevent that right now?

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

41864

Can we just TRUNCATE the Ext nodes without assuming they are extend nodes. That way it just works when you support constants in the future?

Good idea. :) I'll update my patch.

41870

Can we do this without assuming the node is a SIGN/ZERO_EXTEND? Just truncate the original node to Vi8VT.

41870

Can we do this without assuming the node is a SIGN/ZERO_EXTEND? Just truncate the original node to Vi8VT.

Yes, that's better.

42164

This code isn't handling vpdpwssd so why mention it here?

My original code covers both vpdpbusd and vpdpwssd. I'll clean it.

42176

I don't think we can do this if the multiply result is zero extended. Each of the 4 multiplies done by vpdpbusd compute a signed 16-bit product that will be sign extended before adding into the accumulator.

I think we also need to verify that the multiply has at least 2x the number of bits of the input. We shouldn't match (sign_extend (mul (vXi9 (zext (vXi8 X))), (vXi9 (zext (vXi8 Y)))). Does anything prevent that right now?

Really good catch. Thanks.

[LuoYuanke]

Address Craig's comments.

[LuoYuanke]

Align the check in X86ISelLowering.cpp and X86PartialReduction.cpp.
Add test case for 2 x i32.

[LuoYuanke]

Fix lint issues.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

41860

Just name the paramaters ZExt0 and SExt1 and get rid of Ext0 and Ext1?

41870

This just Vi8VT right?

41895

Is this lambda getting anything via & in the capture list or can it just be []

llvm/lib/Target/X86/X86PartialReduction.cpp

99

Why pass 0, nullptr, nullptr when those have default values?

[LuoYuanke]

Address Craig's comments. Thanks, Craig.

[LuoYuanke]

ping

[craig.topper]

LGTM other than that one comment.

llvm/lib/Target/X86/X86ISelLowering.cpp

41862

You can drop these ifs. getZExtOrTrunc/getSExtOrTrunc will do nothing if the type already matches.

[LuoYuanke]

Address Craig's comments and rebase.

llvm/lib/Target/X86/X86ISelLowering.cpp

     SDValue CastedOp0 = DAG.getBitcast(VT, LogicOp0);
     unsigned Opcode = VT.isFloatingPoint() ? FPOpcode : N0.getOpcode();
     return DAG.getNode(Opcode, DL0, VT, LogicOp1.getOperand(0), CastedOp0);
   }
 
   return SDValue();
 }
 
+// (mul (zext a), (sext, b))
+static bool detectExtMul(SelectionDAG &DAG, const SDValue &Mul, SDValue &Op0,
+                         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);
+
+  auto IsFreeTruncation =
+      [](SDValue &Op) -> bool {
+        if ((Op.getOpcode() == ISD::ZERO_EXTEND ||
+             Op.getOpcode() == ISD::SIGN_EXTEND) &&
+            Op.getOperand(0).getValueType().getScalarSizeInBits() <= 8)
+          return true;
+
+        // TODO: Support contant value.
+        return false;
+      };
+
+  auto NumNonSignBits = [&](SDValue &Op) {
+    unsigned NumBits = Op.getValueType().getScalarSizeInBits();
+    unsigned NumSignBits = DAG.ComputeNumSignBits(Op);
+    return NumBits - NumSignBits;
+  };
+
+  // (dpbusd (zext a), (sext, b))
+  if ((IsFreeTruncation(Op0) && NumNonSignBits(Op0) <= 8) &&
+      (IsFreeTruncation(Op1) && NumNonSignBits(Op1) < 8))
+    return true;
+
+  return false;
+}
+
 // Given a ABS node, detect the following pattern:
 // (ABS (SUB (ZERO_EXTEND a), (ZERO_EXTEND b))).
 // This is useful as it is the input into a SAD pattern.
 static bool detectZextAbsDiff(const SDValue &Abs, SDValue &Op0, SDValue &Op1) {
   SDValue AbsOp1 = Abs->getOperand(0);
   if (AbsOp1.getOpcode() != ISD::SUB)
     return false;
 
   Op0 = AbsOp1.getOperand(0);
   Op1 = AbsOp1.getOperand(1);
 
   // Check if the operands of the sub are zero-extended from vectors of i8.
   if (Op0.getOpcode() != ISD::ZERO_EXTEND ||
       Op0.getOperand(0).getValueType().getVectorElementType() != MVT::i8 ||
       Op1.getOpcode() != ISD::ZERO_EXTEND ||
       Op1.getOperand(0).getValueType().getVectorElementType() != MVT::i8)
     return false;
 
   return true;
 }
 
+static SDValue createDP(SelectionDAG &DAG, const SDValue &Ext0,
+                        const SDValue &Ext1, unsigned &LogBias, const SDLoc &DL,
+                        const X86Subtarget &Subtarget) {
+  // Extend or truncate to MVT::i8 first.
+  SDValue ZExt0 = Ext0;
+  SDValue SExt1 = Ext1;
+  MVT Vi8VT =
+      MVT::getVectorVT(MVT::i8, Ext0.getValueType().getVectorElementCount());
+  if (Ext0.getOperand(0).getValueType().getScalarType() != MVT::i8)
+    ZExt0 = DAG.getZExtOrTrunc(Ext0.getOperand(0), DL, Vi8VT);
+  if (Ext1.getOperand(0).getValueType().getScalarType() != MVT::i8)
+    SExt1 = DAG.getSExtOrTrunc(Ext1.getOperand(0), DL, Vi8VT);
+
+  // Find the appropriate width for the DotProduct.
+  EVT InVT = (ZExt0 == Ext0) ? Ext0.getOperand(0).getValueType():
+                               ZExt0.getValueType();
+  LogBias = 2;
+
+  unsigned RegSize = std::max(128u, (unsigned)InVT.getSizeInBits());
+  if (Subtarget.hasVNNI() && !Subtarget.hasVLX())
+    RegSize = std::max(512u, RegSize);
+
+  // "Zero-extend" the i8 vectors. This is not a per-element zext, rather we
+  // fill in the missing vector elements with 0.
+  unsigned NumConcat = RegSize / InVT.getSizeInBits();
+  SmallVector<SDValue, 16> Ops(NumConcat, DAG.getConstant(0, DL, InVT));
+  Ops[0] = (ZExt0 == Ext0) ? Ext0.getOperand(0) : ZExt0;
+  MVT ExtendedVT = MVT::getVectorVT(InVT.getScalarType().getSimpleVT(),
+                                    RegSize / InVT.getScalarSizeInBits());
+  SDValue DpOp0 = DAG.getNode(ISD::CONCAT_VECTORS, DL, ExtendedVT, Ops);
+  Ops[0] = (SExt1 == Ext1) ? Ext1.getOperand(0) : SExt1;
+  SDValue DpOp1 = DAG.getNode(ISD::CONCAT_VECTORS, DL, ExtendedVT, Ops);
+
+  // Actually build the DotProduct, split as 128/256/512 bits for
+  // SSE/AVX2/AVX512BW.
+  auto DpBuilder = [&](SelectionDAG &DAG, const SDLoc &DL,
+                       ArrayRef<SDValue> Ops) {
+    MVT VT = MVT::getVectorVT(MVT::i32, Ops[0].getValueSizeInBits() / 32);
+    return DAG.getNode(X86ISD::VPDPBUSD, DL, VT, Ops);
+  };
+  MVT DpVT = MVT::getVectorVT(MVT::i32, RegSize / 32);
+  SDValue Zero = DAG.getConstant(0, DL, DpVT);
+
+  return SplitOpsAndApply(DAG, Subtarget, DL, DpVT, {Zero, DpOp0, DpOp1},
+                          DpBuilder, false);
+}
+
 // Given two zexts of <k x i8> to <k x i32>, create a PSADBW of the inputs
 // to these zexts.
 static SDValue createPSADBW(SelectionDAG &DAG, const SDValue &Zext0,
                             const SDValue &Zext1, const SDLoc &DL,
                             const X86Subtarget &Subtarget) {
   // Find the appropriate width for the PSADBW.
   EVT InVT = Zext0.getOperand(0).getValueType();
   unsigned RegSize = std::max(128u, (unsigned)InVT.getSizeInBits());
 
   // "Zero-extend" the i8 vectors. This is not a per-element zext, rather we
   // fill in the missing vector elements with 0.
   unsigned NumConcat = RegSize / InVT.getSizeInBits();
   SmallVector<SDValue, 16> Ops(NumConcat, DAG.getConstant(0, DL, InVT));
   Ops[0] = Zext0.getOperand(0);
   MVT ExtendedVT = MVT::getVectorVT(MVT::i8, RegSize / 8);
   SDValue SadOp0 = DAG.getNode(ISD::CONCAT_VECTORS, DL, ExtendedVT, Ops);
   Ops[0] = Zext1.getOperand(0);
   SDValue SadOp1 = DAG.getNode(ISD::CONCAT_VECTORS, DL, ExtendedVT, Ops);
 
   // Actually build the SAD, split as 128/256/512 bits for SSE/AVX2/AVX512BW.
   auto PSADBWBuilder = [](SelectionDAG &DAG, const SDLoc &DL,
                           ArrayRef<SDValue> Ops) {
     MVT VT = MVT::getVectorVT(MVT::i64, Ops[0].getValueSizeInBits() / 64);
     return DAG.getNode(X86ISD::PSADBW, DL, VT, Ops);
   };
   MVT SadVT = MVT::getVectorVT(MVT::i64, RegSize / 64);
   return SplitOpsAndApply(DAG, Subtarget, DL, SadVT, { SadOp0, SadOp1 },
                           PSADBWBuilder);
 }
 
 // Attempt to replace an min/max v8i16/v16i8 horizontal reduction with
 // PHMINPOSUW.
 static SDValue combineMinMaxReduction(SDNode *Extract, SelectionDAG &DAG,
                                       const X86Subtarget &Subtarget) {
   // Bail without SSE41.
   if (!Subtarget.hasSSE41())
     return SDValue();
 
   EVT ExtractVT = Extract->getValueType(0);
   if (ExtractVT != MVT::i16 && ExtractVT != MVT::i8)
     return SDValue();
 
   // Check for SMAX/SMIN/UMAX/UMIN horizontal reduction patterns.
   ISD::NodeType BinOp;
   SDValue Src = DAG.matchBinOpReduction(

   EVT SetccVT =
       TLI.getSetCCResultType(DAG.getDataLayout(), *DAG.getContext(), CmpVT);
   SDValue Setcc = DAG.getSetCC(DL, SetccVT, Movmsk, CmpC, CondCode);
   SDValue Zext = DAG.getZExtOrTrunc(Setcc, DL, ExtractVT);
   SDValue Zero = DAG.getConstant(0, DL, ExtractVT);
   return DAG.getNode(ISD::SUB, DL, ExtractVT, Zero, Zext);
 }
 
+static SDValue combineVNNIPattern(SDNode *Extract, SelectionDAG &DAG,
+                                  const X86Subtarget &Subtarget) {
+  if (!Subtarget.hasVNNI() && !Subtarget.hasAVXVNNI())
+    return SDValue();
+
+  EVT ExtractVT = Extract->getValueType(0);
+  // Verify the type we're extracting is i32, as the output element type of
+  // vpdpbusd and vpdpwssd is i32.
+  if (ExtractVT != MVT::i32)
+    return SDValue();
+
+  EVT VT = Extract->getOperand(0).getValueType();
+  if (!isPowerOf2_32(VT.getVectorNumElements()))
+    return SDValue();
+
+  // Match shuffle + add pyramid.
+  ISD::NodeType BinOp;
+  SDValue Root = DAG.matchBinOpReduction(Extract, BinOp, {ISD::ADD});
+
+  if (Root && (Root.getOpcode() == ISD::SIGN_EXTEND ||
+               Root.getOpcode() == ISD::ZERO_EXTEND ||
+               Root.getOpcode() == ISD::ANY_EXTEND))
+    Root = Root.getOperand(0);
+
+  // If there was a match, we want Root to be a mul.
+  if (!Root || Root.getOpcode() != ISD::MUL)
+    return SDValue();
+
+  // Check whether we have an extend and mul pattern
+  SDValue Ext0, Ext1;
+  if (!detectExtMul(DAG, Root, Ext0, Ext1))
+    return SDValue();
+
+  // Create the dot product instruction.
+  SDLoc DL(Extract);
+  unsigned StageBias;
+  SDValue DP = createDP(DAG, Ext0, Ext1, StageBias, DL, Subtarget);
+
+  // If the original vector was wider than 4 elements, sum over the results
+  // in the DP vector.
+  unsigned Stages = Log2_32(VT.getVectorNumElements());
+  EVT DpVT = DP.getValueType();
+
+  if (Stages > StageBias) {
+    unsigned DpElems = DpVT.getVectorNumElements();
+
+    for(unsigned i = Stages - StageBias; i > 0; --i) {
+      SmallVector<int, 16> Mask(DpElems, -1);
+      for(unsigned j = 0, MaskEnd = 1 << (i - 1); j < MaskEnd; ++j)
+        Mask[j] = MaskEnd + j;
+
+      SDValue Shuffle =
+          DAG.getVectorShuffle(DpVT, DL, DP, DAG.getUNDEF(DpVT), Mask);
+      DP = DAG.getNode(ISD::ADD, DL, DpVT, DP, Shuffle);
+    }
+  }
+
+  unsigned ExtractSizeInBits = ExtractVT.getSizeInBits();
+  // Return the lowest ExtractSizeInBits bits.
+  EVT ResVT = EVT::getVectorVT(*DAG.getContext(), ExtractVT,
+                               DpVT.getSizeInBits() / ExtractSizeInBits);
+  DP = DAG.getBitcast(ResVT, DP);
+  return DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ExtractVT, DP,
+                     Extract->getOperand(1));
+}
+
 static SDValue combineBasicSADPattern(SDNode *Extract, SelectionDAG &DAG,
                                       const X86Subtarget &Subtarget) {
   // PSADBW is only supported on SSE2 and up.
   if (!Subtarget.hasSSE2())
     return SDValue();
 
   EVT ExtractVT = Extract->getValueType(0);
   // Verify the type we're extracting is either i32 or i64.
   // FIXME: Could support other types, but this is what we have coverage for.
   if (ExtractVT != MVT::i32 && ExtractVT != MVT::i64)
     return SDValue();

   if (Root && (Root.getOpcode() == ISD::SIGN_EXTEND ||
                Root.getOpcode() == ISD::ZERO_EXTEND ||
                Root.getOpcode() == ISD::ANY_EXTEND))
     Root = Root.getOperand(0);
 
   // If there was a match, we want Root to be a select that is the root of an
   // abs-diff pattern.
   if (!Root || Root.getOpcode() != ISD::ABS)
     return SDValue();
 
   // Check whether we have an abs-diff pattern feeding into the select.
   SDValue Zext0, Zext1;
   if (!detectZextAbsDiff(Root, Zext0, Zext1))
     return SDValue();
 
   // Create the SAD instruction.
   SDLoc DL(Extract);

   }
 
   // Check whether this extract is the root of a sum of absolute differences
   // pattern. This has to be done here because we really want it to happen
   // pre-legalization,
   if (SDValue SAD = combineBasicSADPattern(N, DAG, Subtarget))
     return SAD;
 
+  if (SDValue VNNI = combineVNNIPattern(N, DAG, Subtarget))
+    return VNNI;
+
   // Attempt to replace an all_of/any_of horizontal reduction with a MOVMSK.
   if (SDValue Cmp = combinePredicateReduction(N, DAG, Subtarget))
     return Cmp;
 
   // Attempt to replace min/max v8i16/v16i8 reductions with PHMINPOSUW.
   if (SDValue MinMax = combineMinMaxReduction(N, DAG, Subtarget))
     return MinMax;
 

llvm/lib/Target/X86/X86PartialReduction.cpp

void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
   }
 
   StringRef getPassName() const override {
     return "X86 Partial Reduction";
   }
 
 private:
-  bool tryMAddReplacement(Instruction *Op);
+  bool tryMAddReplacement(Instruction *Op, bool ReduceInOneBB);
   bool trySADReplacement(Instruction *Op);
 };
 }
 
 FunctionPass *llvm::createX86PartialReductionPass() {
   return new X86PartialReduction();
 }
 
 char X86PartialReduction::ID = 0;
 
 INITIALIZE_PASS(X86PartialReduction, DEBUG_TYPE,
                 "X86 Partial Reduction", false, false)
 
-bool X86PartialReduction::tryMAddReplacement(Instruction *Op) {
+static bool matchVNNIPattern(const X86Subtarget *ST, BinaryOperator *Mul,
+                             const DataLayout *DL) {
+  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);
+
+  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;
+  };
+
+  auto NumNonSignBits = [&](Value *Op) {
+    unsigned NumBits = Mul->getType()->getScalarSizeInBits();
+    unsigned NumSignBits = ComputeNumSignBits(Op, *DL, 0, nullptr, nullptr);
+    return NumBits - NumSignBits;
+  };
+
+  if ((IsFreeTruncation(LHS) && NumNonSignBits(LHS) <= 8) &&
+      (IsFreeTruncation(RHS) && NumNonSignBits(RHS) < 8))
+    return true;
+
+  return false;
+}
+
+bool X86PartialReduction::tryMAddReplacement(Instruction *Op,
+                                             bool ReduceInOneBB) {
   if (!ST->hasSSE2())
     return false;
 
   // Need at least 8 elements.
   if (cast<FixedVectorType>(Op->getType())->getNumElements() < 8)
     return false;
 
   // Element type should be i32.
   if (!cast<VectorType>(Op->getType())->getElementType()->isIntegerTy(32))
     return false;
 
   auto *Mul = dyn_cast<BinaryOperator>(Op);
   if (!Mul || Mul->getOpcode() != Instruction::Mul)
     return false;
 
   Value *LHS = Mul->getOperand(0);
   Value *RHS = Mul->getOperand(1);
 
+  // If the target support VNNI, leave it to ISel to combine reduce operation
+  // to VNNI instruction.
+  // TODO: we can support transforming reduce to VNNI intrinsic for across block
+  // in this pass.
+  if (ReduceInOneBB && matchVNNIPattern(ST, Mul, DL))
+    return false;
+
   // LHS and RHS should be only used once or if they are the same then only
   // used twice. Only check this when SSE4.1 is enabled and we have zext/sext
   // instructions, otherwise we use punpck to emulate zero extend in stages. The
   // trunc/ we need to do likely won't introduce new instructions in that case.
   if (ST->hasSSE41()) {
     if (LHS == RHS) {
       if (!isa<Constant>(LHS) && !LHS->hasNUses(2))
         return false;
bool X86PartialReduction::trySADReplacement(Instruction *Op) {
   SI->replaceAllUsesWith(Ops[0]);
   SI->eraseFromParent();
 
   return true;
 }
 
 // Walk backwards from the ExtractElementInst and determine if it is the end of
 // a horizontal reduction. Return the input to the reduction if we find one.
-static Value *matchAddReduction(const ExtractElementInst &EE) {
+static Value *matchAddReduction(const ExtractElementInst &EE,
+                                bool &ReduceInOneBB) {
+  ReduceInOneBB = true;
   // Make sure we're extracting index 0.
   auto *Index = dyn_cast<ConstantInt>(EE.getIndexOperand());
   if (!Index || !Index->isNullValue())
     return nullptr;
 
   const auto *BO = dyn_cast<BinaryOperator>(EE.getVectorOperand());
   if (!BO || BO->getOpcode() != Instruction::Add || !BO->hasOneUse())
     return nullptr;
+  if (EE.getParent() != BO->getParent())
+    ReduceInOneBB = false;
 
   unsigned NumElems = cast<FixedVectorType>(BO->getType())->getNumElements();
   // Ensure the reduction size is a power of 2.
   if (!isPowerOf2_32(NumElems))
     return nullptr;
 
   const Value *Op = BO;
   unsigned Stages = Log2_32(NumElems);
   for (unsigned i = 0; i != Stages; ++i) {
     const auto *BO = dyn_cast<BinaryOperator>(Op);
     if (!BO || BO->getOpcode() != Instruction::Add)
       return nullptr;
+    if (EE.getParent() != BO->getParent())
+      ReduceInOneBB = false;
 
     // If this isn't the first add, then it should only have 2 users, the
     // shuffle and another add which we checked in the previous iteration.
     if (i != 0 && !BO->hasNUses(2))
       return nullptr;
 
     Value *LHS = BO->getOperand(0);
     Value *RHS = BO->getOperand(1);
bool X86PartialReduction::runOnFunction(Function &F) {
 
   bool MadeChange = false;
   for (auto &BB : F) {
     for (auto &I : BB) {
       auto *EE = dyn_cast<ExtractElementInst>(&I);
       if (!EE)
         continue;
 
+      bool ReduceInOneBB;
       // First find a reduction tree.
       // FIXME: Do we need to handle other opcodes than Add?
-      Value *Root = matchAddReduction(*EE);
+      Value *Root = matchAddReduction(*EE, ReduceInOneBB);
       if (!Root)
         continue;
 
       SmallVector<Instruction *, 8> Leaves;
       collectLeaves(Root, Leaves);
 
       for (Instruction *I : Leaves) {
-        if (tryMAddReplacement(I)) {
+        if (tryMAddReplacement(I, ReduceInOneBB)) {
           MadeChange = true;
           continue;
         }
 
         // Don't do SAD matching on the root node. SelectionDAG already
         // has support for that and currently generates better code.
         if (I != Root && trySADReplacement(I))
           MadeChange = true;
       }
     }
   }
 
   return MadeChange;
 }

llvm/test/CodeGen/X86/dpbusd.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avxvnni | FileCheck %s --check-prefixes=AVXVNNI
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512vnni | FileCheck %s --check-prefixes=AVX512VNNI
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx512vnni -mattr=+avx512vl | FileCheck %s --check-prefixes=AVX512VLVNNI
+
+define dso_local i32 @no_dpbusd(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: no_dpbusd:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vpmovzxbw {{.*#+}} ymm0 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero,mem[4],zero,mem[5],zero,mem[6],zero,mem[7],zero,mem[8],zero,mem[9],zero,mem[10],zero,mem[11],zero,mem[12],zero,mem[13],zero,mem[14],zero,mem[15],zero
+; AVXVNNI-NEXT:    vpmovzxbw {{.*#+}} ymm1 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero,mem[4],zero,mem[5],zero,mem[6],zero,mem[7],zero,mem[8],zero,mem[9],zero,mem[10],zero,mem[11],zero,mem[12],zero,mem[13],zero,mem[14],zero,mem[15],zero
+; AVXVNNI-NEXT:    vpmaddwd %ymm0, %ymm1, %ymm0
+; AVXVNNI-NEXT:    vextracti128 $1, %ymm0, %xmm1
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,2,3]
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vmovd %xmm0, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    vzeroupper
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: no_dpbusd:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vpmovzxbw {{.*#+}} ymm0 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero,mem[4],zero,mem[5],zero,mem[6],zero,mem[7],zero,mem[8],zero,mem[9],zero,mem[10],zero,mem[11],zero,mem[12],zero,mem[13],zero,mem[14],zero,mem[15],zero
+; AVX512VNNI-NEXT:    vpmovzxbw {{.*#+}} ymm1 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero,mem[4],zero,mem[5],zero,mem[6],zero,mem[7],zero,mem[8],zero,mem[9],zero,mem[10],zero,mem[11],zero,mem[12],zero,mem[13],zero,mem[14],zero,mem[15],zero
+; AVX512VNNI-NEXT:    vpmaddwd %ymm0, %ymm1, %ymm0
+; AVX512VNNI-NEXT:    vextracti128 $1, %ymm0, %xmm1
+; AVX512VNNI-NEXT:    vpaddd %ymm1, %ymm0, %ymm0
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,2,3]
+; AVX512VNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: no_dpbusd:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vpmovzxbw {{.*#+}} ymm0 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero,mem[4],zero,mem[5],zero,mem[6],zero,mem[7],zero,mem[8],zero,mem[9],zero,mem[10],zero,mem[11],zero,mem[12],zero,mem[13],zero,mem[14],zero,mem[15],zero
+; AVX512VLVNNI-NEXT:    vpmovzxbw {{.*#+}} ymm1 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero,mem[4],zero,mem[5],zero,mem[6],zero,mem[7],zero,mem[8],zero,mem[9],zero,mem[10],zero,mem[11],zero,mem[12],zero,mem[13],zero,mem[14],zero,mem[15],zero
+; AVX512VLVNNI-NEXT:    vpmaddwd %ymm0, %ymm1, %ymm0
+; AVX512VLVNNI-NEXT:    vextracti128 $1, %ymm0, %xmm1
+; AVX512VLVNNI-NEXT:    vpaddd %ymm1, %ymm0, %ymm0
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[2,3,2,3]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VLVNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    vzeroupper
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0, align 16
+  %2 = zext <16 x i8> %1 to <16 x i32>
+  %3 = bitcast i8* %b to <16 x i8>*
+  %4 = load <16 x i8>, <16 x i8>* %3, align 16
+  %5 = zext <16 x i8> %4 to <16 x i32>
+  %6 = mul nsw <16 x i32> %5, %2
+  %7 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %6)
+  %op.extra = add nsw i32 %7, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @vpdpbusd_mutate(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: vpdpbusd_mutate:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vmovdqa (%rsi), %xmm0
+; AVXVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVXVNNI-NEXT:    {vex} vpdpbusd (%rdi), %xmm0, %xmm1
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vmovd %xmm0, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: vpdpbusd_mutate:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VNNI-NEXT:    vmovdqa (%rsi), %xmm1
+; AVX512VNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VNNI-NEXT:    vpdpbusd %zmm0, %zmm1, %zmm2
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; AVX512VNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: vpdpbusd_mutate:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vmovdqa (%rsi), %xmm0
+; AVX512VLVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512VLVNNI-NEXT:    vpdpbusd (%rdi), %xmm0, %xmm1
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VLVNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0, align 16
+  %2 = sext <16 x i8> %1 to <16 x i32>
+  %3 = bitcast i8* %b to <16 x i8>*
+  %4 = load <16 x i8>, <16 x i8>* %3, align 16
+  %5 = zext <16 x i8> %4 to <16 x i32>
+  %6 = mul nsw <16 x i32> %5, %2
+  %7 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %6)
+  %op.extra = add nsw i32 %7, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @mul_zext(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: mul_zext:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVXVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVXVNNI-NEXT:    {vex} vpdpbusd (%rsi), %xmm0, %xmm1
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vmovd %xmm0, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: mul_zext:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VNNI-NEXT:    vmovdqa (%rsi), %xmm1
+; AVX512VNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VNNI-NEXT:    vpdpbusd %zmm1, %zmm0, %zmm2
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; AVX512VNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: mul_zext:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VLVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512VLVNNI-NEXT:    vpdpbusd (%rsi), %xmm0, %xmm1
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VLVNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0, align 16
+  %2 = zext <16 x i8> %1 to <16 x i16>
+  %3 = bitcast i8* %b to <16 x i8>*
+  %4 = load <16 x i8>, <16 x i8>* %3, align 16
+  %5 = sext <16 x i8> %4 to <16 x i16>
+  %6 = mul nsw <16 x i16> %5, %2
+  %7 = zext <16 x i16> %6 to <16 x i32>
+  %8 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %7)
+  %op.extra = add nsw i32 %8, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @mul_sext(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: mul_sext:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVXVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVXVNNI-NEXT:    {vex} vpdpbusd (%rsi), %xmm0, %xmm1
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vmovd %xmm0, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: mul_sext:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VNNI-NEXT:    vmovdqa (%rsi), %xmm1
+; AVX512VNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VNNI-NEXT:    vpdpbusd %zmm1, %zmm0, %zmm2
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; AVX512VNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: mul_sext:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VLVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512VLVNNI-NEXT:    vpdpbusd (%rsi), %xmm0, %xmm1
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VLVNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0, align 16
+  %2 = zext <16 x i8> %1 to <16 x i16>
+  %3 = bitcast i8* %b to <16 x i8>*
+  %4 = load <16 x i8>, <16 x i8>* %3, align 16
+  %5 = sext <16 x i8> %4 to <16 x i16>
+  %6 = mul nsw <16 x i16> %5, %2
+  %7 = sext <16 x i16> %6 to <16 x i32>
+  %8 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %7)
+  %op.extra = add nsw i32 %8, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @vpdpbusd_512(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: vpdpbusd_512:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVXVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVXVNNI-NEXT:    {vex} vpdpbusd (%rsi), %xmm0, %xmm1
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVXVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT:    vmovd %xmm0, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: vpdpbusd_512:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VNNI-NEXT:    vmovdqa (%rsi), %xmm1
+; AVX512VNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VNNI-NEXT:    vpdpbusd %zmm1, %zmm0, %zmm2
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; AVX512VNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: vpdpbusd_512:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vmovdqa (%rdi), %xmm0
+; AVX512VLVNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512VLVNNI-NEXT:    vpdpbusd (%rsi), %xmm0, %xmm1
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VLVNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0, align 16
+  %2 = zext <16 x i8> %1 to <16 x i32>
+  %3 = bitcast i8* %b to <16 x i8>*
+  %4 = load <16 x i8>, <16 x i8>* %3, align 16
+  %5 = sext <16 x i8> %4 to <16 x i32>
+  %6 = mul nsw <16 x i32> %5, %2
+  %7 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %6)
+  %op.extra = add nsw i32 %7, %c
+  ret i32 %op.extra
+}
+
+declare i32 @llvm.vector.reduce.add.v16i32(<16 x i32>)
+
+define dso_local i32 @vpdpbusd_256(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: vpdpbusd_256:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVXVNNI-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
+; AVXVNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVXVNNI-NEXT:    {vex} vpdpbusd %xmm0, %xmm1, %xmm2
+; AVXVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[1,1,1,1]
+; AVXVNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVXVNNI-NEXT:    vmovd %xmm0, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: vpdpbusd_256:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVX512VNNI-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
+; AVX512VNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VNNI-NEXT:    vpdpbusd %zmm0, %zmm1, %zmm2
+; AVX512VNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[1,1,1,1]
+; AVX512VNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: vpdpbusd_256:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVX512VLVNNI-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
+; AVX512VLVNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VLVNNI-NEXT:    vpdpbusd %xmm0, %xmm1, %xmm2
+; AVX512VLVNNI-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[1,1,1,1]
+; AVX512VLVNNI-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VLVNNI-NEXT:    vmovd %xmm0, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <8 x i8>*
+  %1 = load <8 x i8>, <8 x i8>* %0, align 8
+  %2 = zext <8 x i8> %1 to <8 x i32>
+  %3 = bitcast i8* %b to <8 x i8>*
+  %4 = load <8 x i8>, <8 x i8>* %3, align 8
+  %5 = sext <8 x i8> %4 to <8 x i32>
+  %6 = mul nsw <8 x i32> %5, %2
+  %7 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %6)
+  %op.extra = add nsw i32 %7, %c
+  ret i32 %op.extra
+}
+
+declare i32 @llvm.vector.reduce.add.v8i32(<8 x i32>)
+
+define dso_local i32 @vpdpbusd_128(i8 *%a, i8 *%b, i32 %c, i32 %n) {
+; AVXVNNI-LABEL: vpdpbusd_128:
+; AVXVNNI:       # %bb.0: # %entry
+; AVXVNNI-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVXVNNI-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
+; AVXVNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVXVNNI-NEXT:    vpblendw {{.*#+}} xmm1 = xmm1[0,1],xmm2[2,3,4,5,6,7]
+; AVXVNNI-NEXT:    vpblendw {{.*#+}} xmm0 = xmm0[0,1],xmm2[2,3,4,5,6,7]
+; AVXVNNI-NEXT:    {vex} vpdpbusd %xmm1, %xmm0, %xmm2
+; AVXVNNI-NEXT:    vmovd %xmm2, %eax
+; AVXVNNI-NEXT:    addl %edx, %eax
+; AVXVNNI-NEXT:    retq
+;
+; AVX512VNNI-LABEL: vpdpbusd_128:
+; AVX512VNNI:       # %bb.0: # %entry
+; AVX512VNNI-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVX512VNNI-NEXT:    vpxor %xmm1, %xmm1, %xmm1
+; AVX512VNNI-NEXT:    vpblendw {{.*#+}} xmm0 = xmm0[0,1],xmm1[2,3,4,5,6,7]
+; AVX512VNNI-NEXT:    vmovq {{.*#+}} xmm2 = mem[0],zero
+; AVX512VNNI-NEXT:    vpblendw {{.*#+}} xmm1 = xmm2[0,1],xmm1[2,3,4,5,6,7]
+; AVX512VNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VNNI-NEXT:    vpdpbusd %zmm0, %zmm1, %zmm2
+; AVX512VNNI-NEXT:    vmovd %xmm2, %eax
+; AVX512VNNI-NEXT:    addl %edx, %eax
+; AVX512VNNI-NEXT:    vzeroupper
+; AVX512VNNI-NEXT:    retq
+;
+; AVX512VLVNNI-LABEL: vpdpbusd_128:
+; AVX512VLVNNI:       # %bb.0: # %entry
+; AVX512VLVNNI-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVX512VLVNNI-NEXT:    vmovq {{.*#+}} xmm1 = mem[0],zero
+; AVX512VLVNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VLVNNI-NEXT:    vpblendw {{.*#+}} xmm1 = xmm1[0,1],xmm2[2,3,4,5,6,7]
+; AVX512VLVNNI-NEXT:    vpblendw {{.*#+}} xmm0 = xmm0[0,1],xmm2[2,3,4,5,6,7]
+; AVX512VLVNNI-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; AVX512VLVNNI-NEXT:    vpdpbusd %xmm1, %xmm0, %xmm2
+; AVX512VLVNNI-NEXT:    vmovd %xmm2, %eax
+; AVX512VLVNNI-NEXT:    addl %edx, %eax
+; AVX512VLVNNI-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <4 x i8>*
+  %1 = load <4 x i8>, <4 x i8>* %0, align 8
+  %2 = zext <4 x i8> %1 to <4 x i32>
+  %3 = bitcast i8* %b to <4 x i8>*
+  %4 = load <4 x i8>, <4 x i8>* %3, align 8
+  %5 = sext <4 x i8> %4 to <4 x i32>
+  %6 = mul nsw <4 x i32> %5, %2
+  %7 = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %6)
+  %op.extra = add nsw i32 %7, %c
+  ret i32 %op.extra
+}
+
+declare i32 @llvm.vector.reduce.add.v4i32(<4 x i32>)

llvm/test/CodeGen/X86/dpbusd_i4.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
+
+declare i32 @llvm.vector.reduce.add.v16i32(<16 x i32>)
+
+define dso_local i32 @mul_i8i8(i8 *%a, <16 x i8> %b, i32 %c) {
+; CHECK-LABEL: mul_i8i8:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vmovdqa (%rdi), %xmm1
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm0, %xmm1, %xmm2
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm0, %xmm2, %xmm0
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm1 = xmm0[1,1,1,1]
+; CHECK-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; CHECK-NEXT:    vmovd %xmm0, %eax
+; CHECK-NEXT:    addl %esi, %eax
+; CHECK-NEXT:    retq
+entry:
+  %0 = bitcast i8* %a to <16 x i8>*
+  %1 = load <16 x i8>, <16 x i8>* %0, align 16
+  %2 = zext <16 x i8> %1 to <16 x i32>
+  %3 = sext <16 x i8> %b to <16 x i32>
+  %4 = mul nsw <16 x i32> %2, %3
+  %5 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %4)
+  %op.extra = add nsw i32 %5, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @mul_i4i8(<16 x i4> %a, <16 x i8> %b, i32 %c) {
+; CHECK-LABEL: mul_i4i8:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm1, %xmm0, %xmm2
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm0, %xmm2, %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 i4> %a to <16 x i32>
+  %1 = sext <16 x i8> %b to <16 x i32>
+  %2 = mul nsw <16 x i32> %0, %1
+  %3 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %2)
+  %op.extra = add nsw i32 %3, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @mul_i4i4(<16 x i4> %a, <16 x i4> %b, i32 %c) {
+; CHECK-LABEL: mul_i4i4:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpsllw $4, %xmm1, %xmm1
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm1, %xmm1
+; CHECK-NEXT:    vpsrlw $4, %xmm1, %xmm1
+; CHECK-NEXT:    vmovdqa {{.*#+}} xmm2 = [8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8]
+; CHECK-NEXT:    vpxor %xmm2, %xmm1, %xmm1
+; CHECK-NEXT:    vpsubb %xmm2, %xmm1, %xmm1
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm1, %xmm0, %xmm2
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm0, %xmm2, %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 i4> %a to <16 x i32>
+  %1 = sext <16 x i4> %b to <16 x i32>
+  %2 = mul nsw <16 x i32> %0, %1
+  %3 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %2)
+  %op.extra = add nsw i32 %3, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @mul_sext_i4i4(<16 x i4> %a, <16 x i4> %b, i32 %c) {
+; CHECK-LABEL: mul_sext_i4i4:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpsllw $4, %xmm0, %xmm0
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpsrlw $4, %xmm0, %xmm0
+; CHECK-NEXT:    vmovdqa {{.*#+}} xmm2 = [8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8]
+; CHECK-NEXT:    vpxor %xmm2, %xmm0, %xmm0
+; CHECK-NEXT:    vpsubb %xmm2, %xmm0, %xmm0
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm1, %xmm1
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm0, %xmm1, %xmm2
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm0, %xmm2, %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 <16 x i4> %a to <16 x i32>
+  %1 = sext <16 x i4> %b to <16 x i32>
+  %2 = mul nsw <16 x i32> %0, %1
+  %3 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %2)
+  %op.extra = add nsw i32 %3, %c
+  ret i32 %op.extra
+}
+
+define dso_local i32 @mul_zext_i4i4(<16 x i4> %a, <16 x i4> %b, i32 %c) {
+; CHECK-LABEL: mul_zext_i4i4:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vpsllw $4, %xmm1, %xmm1
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm1, %xmm1
+; CHECK-NEXT:    vpsrlw $4, %xmm1, %xmm1
+; CHECK-NEXT:    vmovdqa {{.*#+}} xmm2 = [8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8]
+; CHECK-NEXT:    vpxor %xmm2, %xmm1, %xmm1
+; CHECK-NEXT:    vpsubb %xmm2, %xmm1, %xmm1
+; CHECK-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; CHECK-NEXT:    vpxor %xmm2, %xmm2, %xmm2
+; CHECK-NEXT:    vpdpbusd %xmm1, %xmm0, %xmm2
+; CHECK-NEXT:    vpshufd {{.*#+}} xmm0 = xmm2[2,3,2,3]
+; CHECK-NEXT:    vpaddd %xmm0, %xmm2, %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 i4> %a to <16 x i32>
+  %1 = zext <16 x i4> %b to <16 x i32>
+  %2 = mul nsw <16 x i32> %0, %1
+  %3 = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> %2)
+  %op.extra = add nsw i32 %3, %c
+  ret i32 %op.extra
+}