[llvm] [X86][LV] Add partial-reduction dot-product support (PR #205373)
Zihao Wang via llvm-commits
llvm-commits at lists.llvm.org
Tue Jun 23 09:29:41 PDT 2026
https://github.com/hhy3 updated https://github.com/llvm/llvm-project/pull/205373
>From 45f50a04a1cea77c5de4ee28f526388906b1c740 Mon Sep 17 00:00:00 2001
From: zh Wang <rekind133 at outlook.com>
Date: Tue, 23 Jun 2026 20:09:42 +0800
Subject: [PATCH] [X86][LV] Add partial-reduction dot-product support
Teach LoopVectorize to recognize reduction chains that can be represented as
llvm.vector.partial.reduce and to account for the narrower accumulator shape
when selecting a VF. Keep the early IR matcher conservative with the VPlan
transform, and only narrow the effective widest type when all widened recurrence
uses are covered.
Add SelectionDAG legalization and expansion support for partial-reduce MLA
forms, plus X86 lowering and costing for VNNI, AVXVNNI, AVXVNNIINT8/16, and BF16
dot-product shapes. On AVX512-register targets that only have 256-bit
dot-product instructions, register the 512-bit forms as Custom so they split to
256-bit operations instead of expanding to zmm arithmetic.
Add LV and CodeGen coverage for constant-trip, runtime-trip, small-VF,
cost-model, and AVX512/VEX split cases.
---
llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h | 1 +
.../SelectionDAG/LegalizeVectorTypes.cpp | 33 +
.../CodeGen/SelectionDAG/TargetLowering.cpp | 4 +
llvm/lib/Target/X86/X86ISelLowering.cpp | 239 +++-
llvm/lib/Target/X86/X86ISelLowering.h | 1 +
.../lib/Target/X86/X86TargetTransformInfo.cpp | 121 ++
llvm/lib/Target/X86/X86TargetTransformInfo.h | 14 +-
.../Vectorize/LoopVectorizationPlanner.cpp | 128 +-
.../Vectorize/LoopVectorizationPlanner.h | 6 +
.../Transforms/Vectorize/VPlanTransforms.cpp | 84 +-
llvm/lib/Transforms/Vectorize/VPlanUtils.cpp | 252 ++++
llvm/lib/Transforms/Vectorize/VPlanUtils.h | 35 +
.../X86/partial-reduce-avx512-vex-split.ll | 89 ++
.../CodeGen/X86/partial-reduce-dot-product.ll | 1033 +++++++++++++++++
.../CodeGen/X86/partial-reduce-small-vf.ll | 38 +
.../X86/partial-reduce-cost-model.ll | 107 ++
.../X86/partial-reduce-dot-product.ll | 653 +++++++++++
.../X86/partial-reduce-small-vf.ll | 96 ++
18 files changed, 2874 insertions(+), 60 deletions(-)
create mode 100644 llvm/test/CodeGen/X86/partial-reduce-avx512-vex-split.ll
create mode 100644 llvm/test/CodeGen/X86/partial-reduce-dot-product.ll
create mode 100644 llvm/test/CodeGen/X86/partial-reduce-small-vf.ll
create mode 100644 llvm/test/Transforms/LoopVectorize/X86/partial-reduce-cost-model.ll
create mode 100644 llvm/test/Transforms/LoopVectorize/X86/partial-reduce-dot-product.ll
create mode 100644 llvm/test/Transforms/LoopVectorize/X86/partial-reduce-small-vf.ll
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index 71d3e1c66be86..6b5f997155d91 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
@@ -1078,6 +1078,7 @@ class LLVM_LIBRARY_VISIBILITY DAGTypeLegalizer {
SDValue WidenVecRes_VECTOR_SHUFFLE(ShuffleVectorSDNode *N);
SDValue WidenVecRes_VECTOR_REVERSE(SDNode *N);
SDValue WidenVecRes_GET_ACTIVE_LANE_MASK(SDNode *N);
+ SDValue WidenVecRes_PARTIAL_REDUCE_MLA(SDNode *N);
void WidenVecRes_VECTOR_DEINTERLEAVE(SDNode *N);
SDValue WidenVecRes_Ternary(SDNode *N);
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index 9e24ae1807ca1..9b782974f0dc8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -5313,6 +5313,12 @@ void DAGTypeLegalizer::WidenVectorResult(SDNode *N, unsigned ResNo) {
case ISD::GET_ACTIVE_LANE_MASK:
Res = WidenVecRes_GET_ACTIVE_LANE_MASK(N);
break;
+ case ISD::PARTIAL_REDUCE_UMLA:
+ case ISD::PARTIAL_REDUCE_SMLA:
+ case ISD::PARTIAL_REDUCE_SUMLA:
+ case ISD::PARTIAL_REDUCE_FMLA:
+ Res = WidenVecRes_PARTIAL_REDUCE_MLA(N);
+ break;
case ISD::VECTOR_DEINTERLEAVE:
WidenVecRes_VECTOR_DEINTERLEAVE(N);
break;
@@ -5601,6 +5607,33 @@ SDValue DAGTypeLegalizer::WidenVecRes_Ternary(SDNode *N) {
{InOp1, InOp2, InOp3, Mask, N->getOperand(4)});
}
+SDValue DAGTypeLegalizer::WidenVecRes_PARTIAL_REDUCE_MLA(SDNode *N) {
+ SDLoc DL(N);
+ EVT AccVT = N->getValueType(0);
+ EVT WidenAccVT = TLI.getTypeToTransformTo(*DAG.getContext(), AccVT);
+ SDValue Acc = GetWidenedVector(N->getOperand(0));
+
+ EVT InputVT = N->getOperand(1).getValueType();
+ ElementCount AccEC = AccVT.getVectorElementCount();
+ ElementCount InputEC = InputVT.getVectorElementCount();
+ assert(InputEC.hasKnownScalarFactor(AccEC) &&
+ "partial-reduce input must be a multiple of the accumulator");
+ unsigned ScaleFactor = InputEC.getKnownScalarFactor(AccEC);
+ EVT WidenInputVT = EVT::getVectorVT(
+ *DAG.getContext(), InputVT.getVectorElementType(),
+ WidenAccVT.getVectorElementCount().multiplyCoefficientBy(ScaleFactor));
+
+ auto WidenInput = [&](SDValue V) {
+ if (getTypeAction(V.getValueType()) == TargetLowering::TypeWidenVector)
+ V = GetWidenedVector(V);
+ return ModifyToType(V, WidenInputVT);
+ };
+
+ return DAG.getNode(N->getOpcode(), DL, WidenAccVT, Acc,
+ WidenInput(N->getOperand(1)),
+ WidenInput(N->getOperand(2)));
+}
+
SDValue DAGTypeLegalizer::WidenVecRes_Binary(SDNode *N) {
// Binary op widening.
SDLoc dl(N);
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 5772ef37ec762..6fae9dfc5856d 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -13770,6 +13770,10 @@ SDValue TargetLowering::expandPartialReduceMLA(SDNode *N,
case ISD::PARTIAL_REDUCE_SMLA:
ExtOpcLHS = ExtOpcRHS = ISD::SIGN_EXTEND;
break;
+ case ISD::PARTIAL_REDUCE_SUMLA:
+ ExtOpcLHS = ISD::SIGN_EXTEND;
+ ExtOpcRHS = ISD::ZERO_EXTEND;
+ break;
case ISD::PARTIAL_REDUCE_FMLA:
ExtOpcLHS = ExtOpcRHS = ISD::FP_EXTEND;
break;
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 291124bb485cb..d2fb8087c5c3e 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -2833,6 +2833,108 @@ X86TargetLowering::X86TargetLowering(const X86TargetMachine &TM,
ISD::INTRINSIC_WO_CHAIN,
ISD::INTRINSIC_W_CHAIN});
+ // Set up partial reduction MLA actions for VNNI dot product support.
+ //
+ // i8 x i8 -> i32:
+ // SUMLA (sext x zext) -> VPDPBUSD (VNNI / AVX-VNNI)
+ // SMLA (sext x sext) -> VPDPBSSD (AVX-VNNI-INT8 / AVX10.2)
+ // UMLA (zext x zext) -> VPDPBUUD (AVX-VNNI-INT8 / AVX10.2)
+ //
+ // i16 x i16 -> i32:
+ // SMLA (sext x sext) -> VPDPWSSD (VNNI / AVX-VNNI)
+ // SUMLA (sext x zext) -> VPDPWSUD (AVX-VNNI-INT16 / AVX10.2)
+ // UMLA (zext x zext) -> VPDPWUUD (AVX-VNNI-INT16 / AVX10.2)
+
+ // VNNI i8: only SUMLA (VPDPBUSD).
+ if (Subtarget.hasVNNI()) {
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v16i32,
+ MVT::v64i8, Custom);
+ if (Subtarget.hasVLX()) {
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v8i32,
+ MVT::v32i8, Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v4i32,
+ MVT::v16i8, Custom);
+ }
+ } else if (Subtarget.hasAVXVNNI()) {
+ if (Subtarget.useAVX512Regs())
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v16i32,
+ MVT::v64i8, Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v8i32, MVT::v32i8,
+ Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SUMLA, MVT::v4i32, MVT::v16i8,
+ Custom);
+ }
+
+ // VNNI-INT8 i8: SMLA (VPDPBSSD), UMLA (VPDPBUUD).
+ if (Subtarget.hasAVX10_2()) {
+ unsigned Int8MLAOps[] = {ISD::PARTIAL_REDUCE_SMLA,
+ ISD::PARTIAL_REDUCE_UMLA};
+ setPartialReduceMLAAction(Int8MLAOps, MVT::v16i32, MVT::v64i8, Custom);
+ setPartialReduceMLAAction(Int8MLAOps, MVT::v8i32, MVT::v32i8, Custom);
+ setPartialReduceMLAAction(Int8MLAOps, MVT::v4i32, MVT::v16i8, Custom);
+ } else if (Subtarget.hasAVXVNNIINT8()) {
+ unsigned Int8MLAOps[] = {ISD::PARTIAL_REDUCE_SMLA,
+ ISD::PARTIAL_REDUCE_UMLA};
+ // On AVX512-register targets, the vectorizer may produce 512-bit
+ // operations.
+ // Register 512-bit as Custom so SplitOpsAndApply can split to 256-bit.
+ if (Subtarget.useAVX512Regs())
+ setPartialReduceMLAAction(Int8MLAOps, MVT::v16i32, MVT::v64i8, Custom);
+ setPartialReduceMLAAction(Int8MLAOps, MVT::v8i32, MVT::v32i8, Custom);
+ setPartialReduceMLAAction(Int8MLAOps, MVT::v4i32, MVT::v16i8, Custom);
+ }
+
+ // VNNI i16: only SMLA (VPDPWSSD).
+ if (Subtarget.hasVNNI()) {
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SMLA, MVT::v16i32,
+ MVT::v32i16, Custom);
+ if (Subtarget.hasVLX()) {
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SMLA, MVT::v8i32,
+ MVT::v16i16, Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SMLA, MVT::v4i32,
+ MVT::v8i16, Custom);
+ }
+ } else if (Subtarget.hasAVXVNNI()) {
+ if (Subtarget.useAVX512Regs())
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SMLA, MVT::v16i32,
+ MVT::v32i16, Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SMLA, MVT::v8i32, MVT::v16i16,
+ Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_SMLA, MVT::v4i32, MVT::v8i16,
+ Custom);
+ }
+
+ // VNNI-INT16 i16: SUMLA (VPDPWSUD), UMLA (VPDPWUUD).
+ if (Subtarget.hasAVX10_2()) {
+ unsigned Int16MLAOps[] = {ISD::PARTIAL_REDUCE_SUMLA,
+ ISD::PARTIAL_REDUCE_UMLA};
+ setPartialReduceMLAAction(Int16MLAOps, MVT::v16i32, MVT::v32i16, Custom);
+ setPartialReduceMLAAction(Int16MLAOps, MVT::v8i32, MVT::v16i16, Custom);
+ setPartialReduceMLAAction(Int16MLAOps, MVT::v4i32, MVT::v8i16, Custom);
+ } else if (Subtarget.hasAVXVNNIINT16()) {
+ unsigned Int16MLAOps[] = {ISD::PARTIAL_REDUCE_SUMLA,
+ ISD::PARTIAL_REDUCE_UMLA};
+ // On AVX512-register targets, the vectorizer may produce 512-bit
+ // operations.
+ // Register 512-bit as Custom so SplitOpsAndApply can split to 256-bit.
+ if (Subtarget.useAVX512Regs())
+ setPartialReduceMLAAction(Int16MLAOps, MVT::v16i32, MVT::v32i16, Custom);
+ setPartialReduceMLAAction(Int16MLAOps, MVT::v8i32, MVT::v16i16, Custom);
+ setPartialReduceMLAAction(Int16MLAOps, MVT::v4i32, MVT::v8i16, Custom);
+ }
+
+ // BF16 dot product: bf16 x bf16 -> f32 (vdpbf16ps, scale=2).
+ if (Subtarget.hasBF16()) {
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v16f32,
+ MVT::v32bf16, Custom);
+ if (Subtarget.hasVLX()) {
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v8f32,
+ MVT::v16bf16, Custom);
+ setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v4f32,
+ MVT::v8bf16, Custom);
+ }
+ }
+
computeRegisterProperties(Subtarget.getRegisterInfo());
MaxStoresPerMemset = 16; // For @llvm.memset -> sequence of stores
@@ -3141,7 +3243,7 @@ static bool isX86CCSigned(X86::CondCode X86CC) {
static X86::CondCode TranslateIntegerX86CC(ISD::CondCode SetCCOpcode) {
switch (SetCCOpcode) {
- // clang-format off
+ // clang-format off
default: llvm_unreachable("Invalid integer condition!");
case ISD::SETEQ: return X86::COND_E;
case ISD::SETGT: return X86::COND_G;
@@ -34570,10 +34672,138 @@ SDValue X86TargetLowering::visitMaskedStore(SelectionDAG &DAG, const SDLoc &DL,
return DAG.getMemIntrinsicNode(X86ISD::CSTORE, DL, Tys, Ops, Ty, MMO);
}
+SDValue X86TargetLowering::LowerPARTIAL_REDUCE_MLA(SDValue Op,
+ SelectionDAG &DAG) const {
+ SDLoc DL(Op);
+ SDValue Acc = Op.getOperand(0);
+ SDValue LHS = Op.getOperand(1);
+ SDValue RHS = Op.getOperand(2);
+ EVT AccVT = Acc.getValueType();
+ unsigned Opcode = Op.getOpcode();
+
+ // BF16 dot product: bf16 x bf16 -> f32 (vdpbf16ps).
+ if (Opcode == ISD::PARTIAL_REDUCE_FMLA) {
+ auto DpBuilder = [](SelectionDAG &DAG, const SDLoc &DL,
+ ArrayRef<SDValue> Ops) {
+ MVT AccVT = MVT::getVectorVT(MVT::f32, Ops[0].getValueSizeInBits() / 32);
+ return DAG.getNode(X86ISD::DPBF16PS, DL, AccVT, Ops);
+ };
+ return SplitOpsAndApply(DAG, Subtarget, DL, AccVT, {Acc, LHS, RHS},
+ DpBuilder, /*CheckBWI=*/false, Subtarget.hasBF16());
+ }
+
+ EVT InputVT = LHS.getValueType();
+ EVT InputEltVT = InputVT.getVectorElementType();
+
+ // Helper to bitcast i16 inputs to i32 (VNNI dot product instructions
+ // operate on i32 register operands that contain packed i16/i8 elements).
+ auto BitcastI16ToI32 = [&](SDValue V) {
+ EVT VT = EVT::getVectorVT(*DAG.getContext(), MVT::i32,
+ V.getValueType().getVectorNumElements() / 2);
+ return DAG.getBitcast(VT, V);
+ };
+
+ bool HasAVX512VNNI = Subtarget.hasVNNI();
+ bool HasVNNIINT8 = Subtarget.hasAVXVNNIINT8() || Subtarget.hasAVX10_2();
+ bool HasVNNIINT16 = Subtarget.hasAVXVNNIINT16() || Subtarget.hasAVX10_2();
+
+ // i16 x i16 -> i32.
+ if (InputEltVT == MVT::i16) {
+ unsigned ISDOpc;
+ SDValue Op1 = LHS, Op2 = RHS;
+ switch (Opcode) {
+ case ISD::PARTIAL_REDUCE_SMLA:
+ // VPDPWSSD: signed x signed (VNNI).
+ ISDOpc = X86ISD::VPDPWSSD;
+ break;
+ case ISD::PARTIAL_REDUCE_SUMLA:
+ // VPDPWSUD: signed x unsigned (VNNI-INT16).
+ // SUMLA: Op1=sext, Op2=zext -> matches VPDPWSUD(acc, signed, unsigned).
+ assert(HasVNNIINT16 && "SUMLA i16 should only be Custom with "
+ "AVXVNNIINT16 or AVX10.2");
+ ISDOpc = X86ISD::VPDPWSUD;
+ break;
+ case ISD::PARTIAL_REDUCE_UMLA:
+ // VPDPWUUD: unsigned x unsigned (VNNI-INT16).
+ assert(HasVNNIINT16 && "UMLA i16 should only be Custom with "
+ "AVXVNNIINT16 or AVX10.2");
+ ISDOpc = X86ISD::VPDPWUUD;
+ break;
+ default:
+ llvm_unreachable("unexpected opcode");
+ }
+
+ SDValue CastLHS = BitcastI16ToI32(Op1);
+ SDValue CastRHS = BitcastI16ToI32(Op2);
+
+ // Allow 512-bit only when the specific instruction has a 512-bit encoding:
+ // VPDPWSSD (SMLA): 512-bit available with AVX512-VNNI.
+ // VPDPWSUD/VPDPWUUD (SUMLA/UMLA): 512-bit available with AVX10.2 only.
+ bool Allow512 = (Opcode == ISD::PARTIAL_REDUCE_SMLA)
+ ? HasAVX512VNNI
+ : Subtarget.hasAVX10_2();
+ auto DpBuilder = [ISDOpc](SelectionDAG &DAG, const SDLoc &DL,
+ ArrayRef<SDValue> Ops) {
+ MVT VT = MVT::getVectorVT(MVT::i32, Ops[0].getValueSizeInBits() / 32);
+ return DAG.getNode(ISDOpc, DL, VT, Ops);
+ };
+ return SplitOpsAndApply(DAG, Subtarget, DL, AccVT, {Acc, CastLHS, CastRHS},
+ DpBuilder,
+ /*CheckBWI=*/false, Allow512);
+ }
+
+ // i8 x i8 -> i32.
+ {
+ unsigned ISDOpc;
+ SDValue Op1, Op2;
+ switch (Opcode) {
+ case ISD::PARTIAL_REDUCE_SUMLA:
+ // VPDPBUSD: unsigned x signed (VNNI).
+ // SUMLA: Op1=sext, Op2=zext -> VPDPBUSD(acc, unsigned, signed).
+ ISDOpc = X86ISD::VPDPBUSD;
+ Op1 = RHS; // zext (unsigned)
+ Op2 = LHS; // sext (signed)
+ break;
+ case ISD::PARTIAL_REDUCE_SMLA:
+ // VPDPBSSD: signed x signed (VNNI-INT8).
+ assert(HasVNNIINT8 && "SMLA i8 should only be Custom with "
+ "AVXVNNIINT8 or AVX10.2");
+ ISDOpc = X86ISD::VPDPBSSD;
+ Op1 = LHS;
+ Op2 = RHS;
+ break;
+ case ISD::PARTIAL_REDUCE_UMLA:
+ // VPDPBUUD: unsigned x unsigned (VNNI-INT8).
+ assert(HasVNNIINT8 && "UMLA i8 should only be Custom with "
+ "AVXVNNIINT8 or AVX10.2");
+ ISDOpc = X86ISD::VPDPBUUD;
+ Op1 = LHS;
+ Op2 = RHS;
+ break;
+ default:
+ llvm_unreachable("unexpected opcode");
+ }
+
+ // Allow 512-bit only when the specific instruction has a 512-bit encoding:
+ // VPDPBUSD (SUMLA): 512-bit available with AVX512-VNNI.
+ // VPDPBSSD/VPDPBUUD (SMLA/UMLA): 512-bit available with AVX10.2 only.
+ bool Allow512 = (Opcode == ISD::PARTIAL_REDUCE_SUMLA)
+ ? HasAVX512VNNI
+ : Subtarget.hasAVX10_2();
+ auto DpBuilder = [ISDOpc](SelectionDAG &DAG, const SDLoc &DL,
+ ArrayRef<SDValue> Ops) {
+ MVT VT = MVT::getVectorVT(MVT::i32, Ops[0].getValueSizeInBits() / 32);
+ return DAG.getNode(ISDOpc, DL, VT, Ops);
+ };
+ return SplitOpsAndApply(DAG, Subtarget, DL, AccVT, {Acc, Op1, Op2},
+ DpBuilder, /*CheckBWI=*/false, Allow512);
+ }
+}
+
/// Provide custom lowering hooks for some operations.
SDValue X86TargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
switch (Op.getOpcode()) {
- // clang-format off
+ // clang-format off
default: llvm_unreachable("Should not custom lower this!");
case ISD::ATOMIC_FENCE: return LowerATOMIC_FENCE(Op, Subtarget, DAG);
case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS:
@@ -34738,6 +34968,11 @@ SDValue X86TargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
case X86ISD::CVTPS2PH: return LowerCVTPS2PH(Op, DAG);
case ISD::PREFETCH: return LowerPREFETCH(Op, Subtarget, DAG);
case ISD::FLDEXP: return LowerFLDEXP(Op, Subtarget, DAG);
+ case ISD::PARTIAL_REDUCE_SMLA:
+ case ISD::PARTIAL_REDUCE_UMLA:
+ case ISD::PARTIAL_REDUCE_SUMLA:
+ case ISD::PARTIAL_REDUCE_FMLA:
+ return LowerPARTIAL_REDUCE_MLA(Op, DAG);
// clang-format on
}
}
diff --git a/llvm/lib/Target/X86/X86ISelLowering.h b/llvm/lib/Target/X86/X86ISelLowering.h
index 0d05c5772a707..b1fdc5a42bf36 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.h
+++ b/llvm/lib/Target/X86/X86ISelLowering.h
@@ -786,6 +786,7 @@ namespace llvm {
SDValue LRINT_LLRINTHelper(SDNode *N, SelectionDAG &DAG) const;
SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
+ SDValue LowerPARTIAL_REDUCE_MLA(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerVSELECT(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
SDValue LowerINSERT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
diff --git a/llvm/lib/Target/X86/X86TargetTransformInfo.cpp b/llvm/lib/Target/X86/X86TargetTransformInfo.cpp
index fba2f08912fed..3388f5979270d 100644
--- a/llvm/lib/Target/X86/X86TargetTransformInfo.cpp
+++ b/llvm/lib/Target/X86/X86TargetTransformInfo.cpp
@@ -5628,6 +5628,127 @@ X86TTIImpl::getAddressComputationCost(Type *PtrTy, ScalarEvolution *SE,
return BaseT::getAddressComputationCost(PtrTy, SE, Ptr, CostKind);
}
+InstructionCost X86TTIImpl::getPartialReductionCost(
+ unsigned Opcode, Type *InputTypeA, Type *InputTypeB, Type *AccumType,
+ ElementCount VF, TTI::PartialReductionExtendKind OpAExtend,
+ TTI::PartialReductionExtendKind OpBExtend, std::optional<unsigned> BinOp,
+ TTI::TargetCostKind CostKind, std::optional<FastMathFlags> FMF) const {
+ InstructionCost Invalid = InstructionCost::getInvalid();
+
+ if (CostKind != TTI::TCK_RecipThroughput)
+ return Invalid;
+
+ if (!BinOp)
+ return Invalid;
+
+ // Both input types must match.
+ if (InputTypeB && InputTypeA != InputTypeB)
+ return Invalid;
+
+ // Floating-point partial reductions require reassoc and contract to allow
+ // the fusion into a single instruction (e.g. vdpbf16ps).
+ if (AccumType->isFloatingPointTy()) {
+ assert(FMF && "Missing FastMathFlags for floating-point partial reduction");
+ if (!FMF->allowReassoc() || !FMF->allowContract())
+ return Invalid;
+ } else {
+ assert(!FMF &&
+ "FastMathFlags only apply to floating-point partial reductions");
+ }
+
+ assert(OpBExtend != TTI::PR_None && InputTypeB &&
+ "Unexpected values for OpBExtend or InputTypeB");
+
+ unsigned ScaleFactor = 0;
+ unsigned PartialReduceOpcode = 0;
+
+ // BF16 dot product: bf16 x bf16 -> f32 (vdpbf16ps, scale=2).
+ if (Opcode == Instruction::FAdd && *BinOp == Instruction::FMul &&
+ InputTypeA->isBFloatTy() && AccumType->isFloatTy()) {
+ if (!ST->hasBF16())
+ return Invalid;
+ ScaleFactor = 2;
+ PartialReduceOpcode = ISD::PARTIAL_REDUCE_FMLA;
+ if (!VF.isKnownMultipleOf(2))
+ return Invalid;
+ }
+ // i8 x i8 -> i32 (vpdpbusd/vpdpbssd/vpdpbuud, scale=4).
+ else if (Opcode == Instruction::Add && *BinOp == Instruction::Mul &&
+ InputTypeA->isIntegerTy(8) && AccumType->isIntegerTy(32)) {
+ ScaleFactor = 4;
+ if (!VF.isKnownMultipleOf(4))
+ return Invalid;
+ if (OpAExtend == TTI::PR_None)
+ return Invalid;
+ // Check that the sign combination is supported by available instructions.
+ bool HasVNNI = ST->hasVNNI() || ST->hasAVXVNNI();
+ bool HasVNNIINT8 = ST->hasAVXVNNIINT8() || ST->hasAVX10_2();
+ bool IsSUMLA = OpAExtend != OpBExtend; // mixed sign -> VPDPBUSD
+ bool IsSMLA =
+ OpAExtend == TTI::PR_SignExtend && OpBExtend == TTI::PR_SignExtend;
+ bool IsUMLA =
+ OpAExtend == TTI::PR_ZeroExtend && OpBExtend == TTI::PR_ZeroExtend;
+ PartialReduceOpcode = IsSUMLA ? ISD::PARTIAL_REDUCE_SUMLA
+ : IsSMLA ? ISD::PARTIAL_REDUCE_SMLA
+ : ISD::PARTIAL_REDUCE_UMLA;
+ if (IsSUMLA && !HasVNNI)
+ return Invalid;
+ if ((IsSMLA || IsUMLA) && !HasVNNIINT8)
+ return Invalid;
+ }
+ // i16 x i16 -> i32 (vpdpwssd/vpdpwsud/vpdpwuud, scale=2).
+ else if (Opcode == Instruction::Add && *BinOp == Instruction::Mul &&
+ InputTypeA->isIntegerTy(16) && AccumType->isIntegerTy(32)) {
+ ScaleFactor = 2;
+ if (!VF.isKnownMultipleOf(2))
+ return Invalid;
+ if (OpAExtend == TTI::PR_None)
+ return Invalid;
+ // Check that the sign combination is supported by available instructions.
+ bool HasVNNI = ST->hasVNNI() || ST->hasAVXVNNI();
+ bool HasVNNIINT16 = ST->hasAVXVNNIINT16() || ST->hasAVX10_2();
+ bool IsSMLA =
+ OpAExtend == TTI::PR_SignExtend && OpBExtend == TTI::PR_SignExtend;
+ bool IsSUMLA = OpAExtend != OpBExtend; // mixed sign -> VPDPWSUD
+ bool IsUMLA =
+ OpAExtend == TTI::PR_ZeroExtend && OpBExtend == TTI::PR_ZeroExtend;
+ PartialReduceOpcode = IsSUMLA ? ISD::PARTIAL_REDUCE_SUMLA
+ : IsSMLA ? ISD::PARTIAL_REDUCE_SMLA
+ : ISD::PARTIAL_REDUCE_UMLA;
+ if (IsSMLA && !HasVNNI)
+ return Invalid;
+ if ((IsSUMLA || IsUMLA) && !HasVNNIINT16)
+ return Invalid;
+ } else {
+ return Invalid;
+ }
+
+ // Check the exact partial-reduction shape has a lowering action. Type
+ // legalization alone may accept smaller accumulator vectors than the X86 dot
+ // product lowerings support.
+ Type *AccVecTy =
+ VectorType::get(AccumType, VF.divideCoefficientBy(ScaleFactor));
+ Type *InputVecTy = VectorType::get(InputTypeA, VF);
+ EVT AccVT = TLI->getValueType(DL, AccVecTy);
+ EVT InputVT = TLI->getValueType(DL, InputVecTy);
+ if (!AccVT.isSimple() || !InputVT.isSimple() ||
+ !TLI->isPartialReduceMLALegalOrCustom(PartialReduceOpcode, AccVT,
+ InputVT))
+ return Invalid;
+
+ auto LT = getTypeLegalizationCost(AccVecTy);
+ if (LT.second == MVT::INVALID_SIMPLE_VALUE_TYPE)
+ return Invalid;
+
+ unsigned CostOpcode =
+ AccumType->isFloatingPointTy() ? Instruction::FMul : Instruction::Mul;
+ InstructionCost DotProductCost =
+ getArithmeticInstrCost(CostOpcode, AccVecTy, CostKind);
+ if (!DotProductCost.isValid())
+ return Invalid;
+ return std::max(InstructionCost(LT.first), DotProductCost);
+}
+
InstructionCost
X86TTIImpl::getArithmeticReductionCost(unsigned Opcode, VectorType *ValTy,
std::optional<FastMathFlags> FMF,
diff --git a/llvm/lib/Target/X86/X86TargetTransformInfo.h b/llvm/lib/Target/X86/X86TargetTransformInfo.h
index 4120421622b21..e77a1f6dadeac 100644
--- a/llvm/lib/Target/X86/X86TargetTransformInfo.h
+++ b/llvm/lib/Target/X86/X86TargetTransformInfo.h
@@ -222,19 +222,17 @@ class X86TTIImpl final : public BasicTTIImplBase<X86TTIImpl> {
getIntrinsicInstrCost(const IntrinsicCostAttributes &ICA,
TTI::TargetCostKind CostKind) const override;
- InstructionCost
- getArithmeticReductionCost(unsigned Opcode, VectorType *Ty,
- std::optional<FastMathFlags> FMF,
- TTI::TargetCostKind CostKind) const override;
-
InstructionCost getPartialReductionCost(
unsigned Opcode, Type *InputTypeA, Type *InputTypeB, Type *AccumType,
ElementCount VF, TTI::PartialReductionExtendKind OpAExtend,
TTI::PartialReductionExtendKind OpBExtend, std::optional<unsigned> BinOp,
TTI::TargetCostKind CostKind,
- std::optional<FastMathFlags> FMF) const override {
- return InstructionCost::getInvalid();
- }
+ std::optional<FastMathFlags> FMF) const override;
+
+ InstructionCost
+ getArithmeticReductionCost(unsigned Opcode, VectorType *Ty,
+ std::optional<FastMathFlags> FMF,
+ TTI::TargetCostKind CostKind) const override;
InstructionCost getMinMaxCost(Intrinsic::ID IID, Type *Ty,
TTI::TargetCostKind CostKind,
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.cpp
index 509fe991d2652..fe5f979b0a835 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.cpp
@@ -13,6 +13,8 @@
//===----------------------------------------------------------------------===//
#include "LoopVectorizationPlanner.h"
+#include "VPlanUtils.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
#include "llvm/Analysis/ScalarEvolution.h"
@@ -397,16 +399,124 @@ FixedScalableVFPair VFSelectionContext::computeFeasibleMaxVF(
bool FoldTailByMasking, bool RequiresScalarEpilogue) {
auto [SmallestType, WidestType] = getSmallestAndWidestTypes();
+ // Save original WidestType for dependence distance safety calculation.
+ // Partial reduction narrowing affects VF selection but must not inflate
+ // MaxSafeElementsPowerOf2, which depends on actual memory access widths.
+ unsigned WidestTypeForSafety = WidestType;
+
+ // Check if partial reductions allow narrowing the effective WidestType. This
+ // is only valid for recurrence types whose loop-wide uses all come from
+ // reductions that can be represented as partial reductions. Non-reduction
+ // loads/stores of the same type must keep the original type width.
+ if (!ElementTypesInLoop.empty()) {
+ const DataLayout &DL = F.getDataLayout();
+ DenseMap<Type *, unsigned> PartialReductionWidths;
+ SmallPtrSet<Type *, 4> UnsupportedReductionTypes;
+
+ for (const auto &[Phi, RdxDesc] : Legal->getReductionVars()) {
+ Type *AccumType = RdxDesc.getRecurrenceType();
+ if (!ReductionElementTypesInLoop.contains(AccumType))
+ continue;
+
+ auto PartialReductionChain =
+ vputils::getPartialReductionChain(Phi, RdxDesc, TheLoop);
+ if (!PartialReductionChain) {
+ UnsupportedReductionTypes.insert(AccumType);
+ continue;
+ }
+
+ unsigned WidestInputType = 0;
+ bool Supported = true;
+ for (const vputils::PartialReductionChainLink &Link :
+ *PartialReductionChain) {
+ unsigned InputAWidth = Link.InputA.SrcType->getScalarSizeInBits();
+ unsigned InputBWidth = Link.InputB.SrcType
+ ? Link.InputB.SrcType->getScalarSizeInBits()
+ : InputAWidth;
+ if (!InputAWidth || !InputBWidth) {
+ Supported = false;
+ break;
+ }
+
+ // Start with the smallest potentially useful VF, but keep probing up
+ // to one register of accumulators. Some targets only support partial
+ // reductions once the accumulator has enough lanes.
+ unsigned MinProbeVF = Link.ScaleFactor * Link.ScaleFactor;
+ unsigned MaxProbeVF = MinProbeVF;
+ unsigned AccumWidth = AccumType->getScalarSizeInBits();
+ TypeSize RegisterWidth =
+ TTI.getRegisterBitWidth(TTI::RGK_FixedWidthVector);
+ if (AccumWidth && !RegisterWidth.isScalable()) {
+ unsigned RegisterProbeVF =
+ Link.ScaleFactor *
+ static_cast<unsigned>(
+ divideCeil(RegisterWidth.getFixedValue(), AccumWidth));
+ MaxProbeVF = std::max(MaxProbeVF, RegisterProbeVF);
+ }
+
+ bool LinkSupported = false;
+ for (unsigned ProbeVF = MinProbeVF; ProbeVF <= MaxProbeVF;
+ ProbeVF += Link.ScaleFactor) {
+ if (vputils::getPartialReductionCost(
+ Link, TTI, ElementCount::getFixed(ProbeVF), CostKind)
+ .isValid()) {
+ LinkSupported = true;
+ break;
+ }
+ }
+
+ if (!LinkSupported) {
+ Supported = false;
+ break;
+ }
+ WidestInputType =
+ std::max(WidestInputType, std::max(InputAWidth, InputBWidth));
+ }
+
+ if (!Supported || !WidestInputType ||
+ WidestInputType >= AccumType->getScalarSizeInBits()) {
+ UnsupportedReductionTypes.insert(AccumType);
+ continue;
+ }
+
+ auto [It, Inserted] =
+ PartialReductionWidths.try_emplace(AccumType, WidestInputType);
+ if (!Inserted)
+ It->second = std::max(It->second, WidestInputType);
+ }
+
+ for (Type *Ty : UnsupportedReductionTypes)
+ PartialReductionWidths.erase(Ty);
+
+ unsigned EffectiveWidestType = 8;
+ for (Type *Ty : ElementTypesInLoop) {
+ unsigned TypeWidth =
+ DL.getTypeSizeInBits(Ty->getScalarType()).getFixedValue();
+ unsigned EffectiveTypeWidth = TypeWidth;
+ if (!NonReductionElementTypesInLoop.contains(Ty))
+ if (auto It = PartialReductionWidths.find(Ty);
+ It != PartialReductionWidths.end())
+ EffectiveTypeWidth = It->second;
+ EffectiveWidestType = std::max(EffectiveWidestType, EffectiveTypeWidth);
+ }
+
+ if (EffectiveWidestType < WidestType) {
+ LLVM_DEBUG(dbgs() << "LV: Partial reductions narrow WidestType from "
+ << WidestType << " to " << EffectiveWidestType << "\n");
+ WidestType = EffectiveWidestType;
+ }
+ }
+
// Get the maximum safe dependence distance in bits computed by LAA.
// It is computed by MaxVF * sizeOf(type) * 8, where type is taken from
// the memory accesses that is most restrictive (involved in the smallest
// dependence distance).
- unsigned MaxSafeElementsPowerOf2 =
- llvm::bit_floor(Legal->getMaxSafeVectorWidthInBits() / WidestType);
+ unsigned MaxSafeElementsPowerOf2 = llvm::bit_floor(
+ Legal->getMaxSafeVectorWidthInBits() / WidestTypeForSafety);
if (!Legal->isSafeForAnyStoreLoadForwardDistances()) {
unsigned SLDist = Legal->getMaxStoreLoadForwardSafeDistanceInBits();
MaxSafeElementsPowerOf2 =
- std::min(MaxSafeElementsPowerOf2, SLDist / WidestType);
+ std::min(MaxSafeElementsPowerOf2, SLDist / WidestTypeForSafety);
}
auto MaxSafeFixedVF = ElementCount::getFixed(MaxSafeElementsPowerOf2);
@@ -537,11 +647,14 @@ VFSelectionContext::getSmallestAndWidestTypes() const {
void VFSelectionContext::collectElementTypesForWidening(
const SmallPtrSetImpl<const Value *> *ValuesToIgnore) {
ElementTypesInLoop.clear();
+ NonReductionElementTypesInLoop.clear();
+ ReductionElementTypesInLoop.clear();
// For each block.
for (BasicBlock *BB : TheLoop->blocks()) {
// For each instruction in the loop.
for (Instruction &I : *BB) {
Type *T = I.getType();
+ bool IsReductionElementType = false;
// Skip ignored values.
if (ValuesToIgnore && ValuesToIgnore->contains(&I))
@@ -563,16 +676,23 @@ void VFSelectionContext::collectElementTypesForWidening(
RdxDesc.getRecurrenceType()))
continue;
T = RdxDesc.getRecurrenceType();
+ IsReductionElementType = true;
}
// Examine the stored values.
- if (auto *ST = dyn_cast<StoreInst>(&I))
+ if (auto *ST = dyn_cast<StoreInst>(&I)) {
T = ST->getValueOperand()->getType();
+ IsReductionElementType = false;
+ }
assert(T->isSized() &&
"Expected the load/store/recurrence type to be sized");
ElementTypesInLoop.insert(T);
+ if (IsReductionElementType)
+ ReductionElementTypesInLoop.insert(T);
+ else
+ NonReductionElementTypesInLoop.insert(T);
}
}
}
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
index b39ffdebf6179..d92244c91f204 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
@@ -663,6 +663,12 @@ class VFSelectionContext {
/// All element types found in the loop.
SmallPtrSet<Type *, 16> ElementTypesInLoop;
+ /// Element types found in the loop via values other than reduction PHIs.
+ SmallPtrSet<Type *, 16> NonReductionElementTypesInLoop;
+
+ /// Element types found in the loop via reduction PHIs.
+ SmallPtrSet<Type *, 16> ReductionElementTypesInLoop;
+
/// PHINodes of the reductions that should be expanded in-loop. Set by
/// collectInLoopReductions.
SmallPtrSet<PHINode *, 4> InLoopReductions;
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index adcfe30ff9561..8cc5050da94ce 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -307,28 +307,28 @@ static bool sinkScalarOperands(VPlan &Plan) {
bool Changed = false;
SetVector<std::pair<VPBasicBlock *, VPSingleDefRecipe *>> WorkList;
- auto InsertIfValidSinkCandidate = [ScalarVFOnly, &WorkList](
- VPBasicBlock *SinkTo, VPValue *Op) {
- auto *Candidate =
- dyn_cast_or_null<VPSingleDefRecipe>(Op->getDefiningRecipe());
- if (!Candidate)
- return;
-
- // We only know how to sink VPReplicateRecipes and VPScalarIVStepsRecipes
- // for now.
- if (!isa<VPReplicateRecipe, VPScalarIVStepsRecipe>(Candidate))
- return;
-
- if (Candidate->getParent() == SinkTo ||
- vputils::cannotHoistOrSinkRecipe(*Candidate, /*Sinking=*/true))
- return;
-
- if (auto *RepR = dyn_cast<VPReplicateRecipe>(Candidate))
- if (!ScalarVFOnly && RepR->isSingleScalar())
- return;
-
- WorkList.insert({SinkTo, Candidate});
- };
+ auto InsertIfValidSinkCandidate =
+ [ScalarVFOnly, &WorkList](VPBasicBlock *SinkTo, VPValue *Op) {
+ auto *Candidate =
+ dyn_cast_or_null<VPSingleDefRecipe>(Op->getDefiningRecipe());
+ if (!Candidate)
+ return;
+
+ // We only know how to sink VPReplicateRecipes and
+ // VPScalarIVStepsRecipes for now.
+ if (!isa<VPReplicateRecipe, VPScalarIVStepsRecipe>(Candidate))
+ return;
+
+ if (Candidate->getParent() == SinkTo ||
+ vputils::cannotHoistOrSinkRecipe(*Candidate, /*Sinking=*/true))
+ return;
+
+ if (auto *RepR = dyn_cast<VPReplicateRecipe>(Candidate))
+ if (!ScalarVFOnly && RepR->isSingleScalar())
+ return;
+
+ WorkList.insert({SinkTo, Candidate});
+ };
// First, collect the operands of all recipes in replicate blocks as seeds for
// sinking.
@@ -6809,30 +6809,22 @@ getPartialReductionLinkCost(VPCostContext &CostCtx,
ElementCount VF) {
Type *RdxType = Link.ReductionBinOp->getScalarType();
const ExtendedReductionOperand &ExtendedOp = Link.ExtendedOp;
- std::optional<unsigned> BinOpc = std::nullopt;
- // If ExtendB is not none, then the "ExtendsUser" is the binary operation.
- if (ExtendedOp.ExtendB.Kind != ExtendKind::PR_None)
- BinOpc = ExtendedOp.ExtendsUser->getOpcode();
-
- std::optional<llvm::FastMathFlags> Flags;
- if (RdxType->isFloatingPointTy())
- Flags = Link.ReductionBinOp->getFastMathFlagsOrNone();
-
- auto GetLinkOpcode = [&Link]() -> unsigned {
- switch (Link.RK) {
- case RecurKind::Sub:
- return Instruction::Add;
- case RecurKind::FSub:
- return Instruction::FAdd;
- default:
- return Link.ReductionBinOp->getOpcode();
- }
- };
-
- return CostCtx.TTI.getPartialReductionCost(
- GetLinkOpcode(), ExtendedOp.ExtendA.SrcType, ExtendedOp.ExtendB.SrcType,
- RdxType, VF, ExtendedOp.ExtendA.Kind, ExtendedOp.ExtendB.Kind, BinOpc,
- CostCtx.CostKind, Flags);
+ vputils::PartialReductionChainLink CostLink{
+ /*AccumType=*/RdxType,
+ /*InputA=*/{ExtendedOp.ExtendA.SrcType, ExtendedOp.ExtendA.Kind},
+ /*InputB=*/{ExtendedOp.ExtendB.SrcType, ExtendedOp.ExtendB.Kind},
+ /*BinOp=*/ExtendedOp.ExtendB.Kind != ExtendKind::PR_None
+ ? std::optional<unsigned>(ExtendedOp.ExtendsUser->getOpcode())
+ : std::nullopt,
+ /*FMF=*/RdxType->isFloatingPointTy()
+ ? std::optional<FastMathFlags>(
+ Link.ReductionBinOp->getFastMathFlagsOrNone())
+ : std::nullopt,
+ /*RK=*/Link.RK,
+ /*ReductionOpcode=*/Link.ReductionBinOp->getOpcode(),
+ /*ScaleFactor=*/Link.ScaleFactor};
+ return vputils::getPartialReductionCost(CostLink, CostCtx.TTI, VF,
+ CostCtx.CostKind);
}
static ExtendKind getPartialReductionExtendKind(VPWidenCastRecipe *Cast) {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp b/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
index 7e0f536a63a83..1b924d3711590 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
@@ -11,18 +11,178 @@
#include "VPlanAnalysis.h"
#include "VPlanCFG.h"
#include "VPlanDominatorTree.h"
+#include "VPlanHelpers.h"
#include "VPlanPatternMatch.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/MemoryLocation.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/ScalarEvolutionPatternMatch.h"
+#include "llvm/Analysis/VectorUtils.h"
#include "llvm/IR/Dominators.h"
+#include "llvm/IR/PatternMatch.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
using namespace llvm;
using namespace llvm::VPlanPatternMatch;
using namespace llvm::SCEVPatternMatch;
+namespace {
+
+using ExtendKind = TargetTransformInfo::PartialReductionExtendKind;
+using PartialReductionChainLink = vputils::PartialReductionChainLink;
+using PartialReductionInput = vputils::PartialReductionInput;
+namespace IRPatternMatch = llvm::PatternMatch;
+
+struct ExtendedReductionOperand {
+ Instruction *ExtendsUser = nullptr;
+ SmallVector<Instruction *, 2> ExtendInsts;
+ PartialReductionInput InputA;
+ PartialReductionInput InputB;
+ std::optional<unsigned> BinOp;
+};
+
+// This IR-level matcher is the conservative pre-VPlan counterpart of the VPlan
+// recipe matcher used by VPlanTransforms::createPartialReductions(). It probes
+// whether partial reductions can affect VF limits, so it must not recognize
+// patterns that the VPlan transform cannot also form. Keep the two matchers in
+// sync when adding a new partial-reduction shape, with paired transform and
+// codegen coverage.
+static std::optional<ExtendedReductionOperand>
+matchExtendedReductionOperand(Instruction *UpdateI, Value *Op) {
+ std::optional<ExtendKind> OuterExtKind;
+ if (auto *Cast = dyn_cast<CastInst>(Op)) {
+ ExtendKind CastKind =
+ TargetTransformInfo::getPartialReductionExtendKind(Cast);
+ if (CastKind == ExtendKind::PR_None)
+ return std::nullopt;
+
+ Value *CastSource = Cast->getOperand(0);
+ OuterExtKind = CastKind;
+ if (IRPatternMatch::match(
+ CastSource, IRPatternMatch::m_Mul(IRPatternMatch::m_Value(),
+ IRPatternMatch::m_Value())) ||
+ IRPatternMatch::match(
+ CastSource, IRPatternMatch::m_FMul(IRPatternMatch::m_Value(),
+ IRPatternMatch::m_Value())))
+ Op = CastSource;
+ else if (UpdateI->getOpcode() == Instruction::Add ||
+ UpdateI->getOpcode() == Instruction::FAdd)
+ return ExtendedReductionOperand{/*ExtendsUser=*/UpdateI,
+ /*ExtendInsts=*/{Cast},
+ /*InputA=*/{Cast->getSrcTy(), CastKind},
+ /*InputB=*/{},
+ /*BinOp=*/std::nullopt};
+ }
+
+ if (auto *I = dyn_cast<Instruction>(Op); I && !I->hasOneUse())
+ return std::nullopt;
+
+ auto *BinOp = dyn_cast<BinaryOperator>(Op);
+ if (!BinOp)
+ return std::nullopt;
+
+ auto *LHSCast = dyn_cast<CastInst>(BinOp->getOperand(0));
+ if (!LHSCast)
+ return std::nullopt;
+
+ ExtendKind LHSExtendKind =
+ TargetTransformInfo::getPartialReductionExtendKind(LHSCast);
+ if (LHSExtendKind == ExtendKind::PR_None)
+ return std::nullopt;
+
+ auto KindsMatchOuterExtend = [&](CastInst *Cast) {
+ return !OuterExtKind || TargetTransformInfo::getPartialReductionExtendKind(
+ Cast) == *OuterExtKind;
+ };
+ if (!KindsMatchOuterExtend(LHSCast))
+ return std::nullopt;
+
+ Type *RHSInputType = LHSCast->getSrcTy();
+ ExtendKind RHSExtendKind = LHSExtendKind;
+ CastInst *RHSCast = nullptr;
+
+ if ((RHSCast = dyn_cast<CastInst>(BinOp->getOperand(1)))) {
+ RHSExtendKind = TargetTransformInfo::getPartialReductionExtendKind(RHSCast);
+ if (RHSExtendKind == ExtendKind::PR_None || !KindsMatchOuterExtend(RHSCast))
+ return std::nullopt;
+ RHSInputType = RHSCast->getSrcTy();
+ } else {
+ const APInt *RHSConst = nullptr;
+ if (!IRPatternMatch::match(BinOp->getOperand(1),
+ IRPatternMatch::m_APInt(RHSConst)) ||
+ !llvm::canConstantBeExtended(RHSConst, LHSCast->getSrcTy(),
+ LHSExtendKind))
+ return std::nullopt;
+ }
+
+ SmallVector<Instruction *, 2> ExtendInsts;
+ ExtendInsts.push_back(LHSCast);
+ if (RHSCast)
+ ExtendInsts.push_back(RHSCast);
+
+ return ExtendedReductionOperand{
+ /*ExtendsUser=*/BinOp,
+ /*ExtendInsts=*/std::move(ExtendInsts),
+ /*InputA=*/{LHSCast->getSrcTy(), LHSExtendKind},
+ /*InputB=*/{RHSInputType, RHSExtendKind},
+ /*BinOp=*/BinOp->getOpcode()};
+}
+
+struct IRPartialReductionChainLink {
+ PartialReductionChainLink CostLink;
+ Instruction *ReductionInst = nullptr;
+ Instruction *ExtendsUser = nullptr;
+ SmallVector<Instruction *, 2> ExtendInsts;
+};
+
+static bool
+isValidPartialReductionChain(ArrayRef<IRPartialReductionChainLink> Chain,
+ PHINode *Phi, const RecurrenceDescriptor &RdxDesc,
+ const Loop *L) {
+ if (Chain.empty() || RdxDesc.IntermediateStore)
+ return false;
+
+ SmallPtrSet<Instruction *, 4> PartialReductionOps;
+ DenseMap<Instruction *, unsigned> ScaledReductionMap;
+ for (const IRPartialReductionChainLink &Link : Chain) {
+ PartialReductionOps.insert(Link.ExtendsUser);
+ ScaledReductionMap[Link.ReductionInst] = Link.CostLink.ScaleFactor;
+ }
+
+ auto ExtendUsersValid = [&](Instruction *Ext) {
+ return all_of(Ext->users(), [&](User *U) {
+ auto *I = dyn_cast<Instruction>(U);
+ return I && PartialReductionOps.contains(I);
+ });
+ };
+
+ for (const IRPartialReductionChainLink &Link : Chain) {
+ if (!all_of(Link.ExtendInsts, ExtendUsersValid))
+ return false;
+
+ auto UseIsValid = [&](User *U) {
+ if (U == Phi)
+ return true;
+
+ auto *I = dyn_cast<Instruction>(U);
+ if (!I)
+ return false;
+
+ if (L->contains(I->getParent()))
+ return Link.CostLink.ScaleFactor == ScaledReductionMap.lookup(I);
+
+ return Link.ReductionInst == RdxDesc.getLoopExitInstr();
+ };
+ if (!all_of(Link.ReductionInst->users(), UseIsValid))
+ return false;
+ }
+
+ return true;
+}
+
+} // namespace
+
bool vputils::onlyFirstLaneUsed(const VPValue *Def) {
return all_of(Def->users(),
[Def](const VPUser *U) { return U->usesFirstLaneOnly(Def); });
@@ -749,6 +909,98 @@ VPInstruction *vputils::findComputeReductionResult(VPReductionPHIRecipe *PhiR) {
cast<VPSingleDefRecipe>(SelR));
}
+std::optional<SmallVector<vputils::PartialReductionChainLink>>
+vputils::getPartialReductionChain(PHINode *Phi,
+ const RecurrenceDescriptor &RdxDesc,
+ const Loop *L) {
+ RecurKind RK = RdxDesc.getRecurrenceKind();
+ if (RK != RecurKind::Add && RK != RecurKind::FAdd && RK != RecurKind::Sub &&
+ RK != RecurKind::FSub)
+ return std::nullopt;
+
+ Type *AccumType = RdxDesc.getRecurrenceType();
+ unsigned AccumWidth = AccumType->getScalarSizeInBits();
+ if (!AccumWidth)
+ return std::nullopt;
+
+ SmallVector<PartialReductionChainLink> Chain;
+ SmallVector<IRPartialReductionChainLink> IRChain;
+ Value *CurrentValue = RdxDesc.getLoopExitInstr();
+ while (CurrentValue != Phi) {
+ auto *UpdateI = dyn_cast<Instruction>(CurrentValue);
+ if (!UpdateI || !L->contains(UpdateI->getParent()) ||
+ !Instruction::isBinaryOp(UpdateI->getOpcode()))
+ return std::nullopt;
+
+ Value *Op = UpdateI->getOperand(1);
+ Value *PrevValue = UpdateI->getOperand(0);
+ std::optional<ExtendedReductionOperand> ExtendedOp =
+ matchExtendedReductionOperand(UpdateI, Op);
+ if (!ExtendedOp) {
+ if (RK == RecurKind::Sub || RK == RecurKind::FSub)
+ return std::nullopt;
+ ExtendedOp = matchExtendedReductionOperand(UpdateI, PrevValue);
+ if (!ExtendedOp)
+ return std::nullopt;
+ std::swap(Op, PrevValue);
+ }
+
+ unsigned InputWidth = ExtendedOp->InputA.SrcType->getScalarSizeInBits();
+ if (!InputWidth || AccumWidth % InputWidth != 0)
+ return std::nullopt;
+ unsigned ScaleFactor = AccumWidth / InputWidth;
+ if (ScaleFactor <= 1)
+ return std::nullopt;
+
+ std::optional<FastMathFlags> FMF;
+ if (AccumType->isFloatingPointTy())
+ FMF = cast<FPMathOperator>(UpdateI)->getFastMathFlags();
+
+ IRChain.push_back({/*CostLink=*/
+ {/*AccumType=*/AccumType,
+ /*InputA=*/ExtendedOp->InputA,
+ /*InputB=*/ExtendedOp->InputB,
+ /*BinOp=*/ExtendedOp->BinOp,
+ /*FMF=*/FMF,
+ /*RK=*/RK,
+ /*ReductionOpcode=*/UpdateI->getOpcode(),
+ /*ScaleFactor=*/ScaleFactor},
+ /*ReductionInst=*/UpdateI,
+ /*ExtendsUser=*/ExtendedOp->ExtendsUser,
+ /*ExtendInsts=*/ExtendedOp->ExtendInsts});
+ CurrentValue = PrevValue;
+ }
+
+ std::reverse(IRChain.begin(), IRChain.end());
+ if (!isValidPartialReductionChain(IRChain, Phi, RdxDesc, L))
+ return std::nullopt;
+
+ Chain.reserve(IRChain.size());
+ for (const IRPartialReductionChainLink &Link : IRChain)
+ Chain.push_back(Link.CostLink);
+ return Chain;
+}
+
+InstructionCost vputils::getPartialReductionCost(
+ const PartialReductionChainLink &Link, const TargetTransformInfo &TTI,
+ ElementCount VF, TTI::TargetCostKind CostKind) {
+ unsigned Opcode;
+ switch (Link.RK) {
+ case RecurKind::Sub:
+ Opcode = Instruction::Add;
+ break;
+ case RecurKind::FSub:
+ Opcode = Instruction::FAdd;
+ break;
+ default:
+ Opcode = Link.ReductionOpcode;
+ break;
+ }
+ return TTI.getPartialReductionCost(
+ Opcode, Link.InputA.SrcType, Link.InputB.SrcType, Link.AccumType, VF,
+ Link.InputA.Kind, Link.InputB.Kind, Link.BinOp, CostKind, Link.FMF);
+}
+
bool vputils::isUsedByLoadStoreAddress(const VPValue *V) {
SmallPtrSet<const VPValue *, 4> Seen;
SmallVector<const VPValue *> WorkList = {V};
diff --git a/llvm/lib/Transforms/Vectorize/VPlanUtils.h b/llvm/lib/Transforms/Vectorize/VPlanUtils.h
index 3ebf8e63d5b36..5d3e8d8f40fdd 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanUtils.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanUtils.h
@@ -10,11 +10,15 @@
#define LLVM_TRANSFORMS_VECTORIZE_VPLANUTILS_H
#include "VPlan.h"
+#include "llvm/Analysis/IVDescriptors.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Support/Compiler.h"
namespace llvm {
class DominatorTree;
class MemoryLocation;
+class Loop;
+class PHINode;
class ScalarEvolution;
class SCEV;
class PredicatedScalarEvolution;
@@ -23,6 +27,24 @@ class PredicatedScalarEvolution;
namespace llvm {
namespace vputils {
+
+struct PartialReductionInput {
+ Type *SrcType = nullptr;
+ TargetTransformInfo::PartialReductionExtendKind Kind =
+ TargetTransformInfo::PR_None;
+};
+
+struct PartialReductionChainLink {
+ Type *AccumType = nullptr;
+ PartialReductionInput InputA;
+ PartialReductionInput InputB;
+ std::optional<unsigned> BinOp;
+ std::optional<FastMathFlags> FMF;
+ RecurKind RK = RecurKind::None;
+ unsigned ReductionOpcode = 0;
+ unsigned ScaleFactor = 0;
+};
+
/// Returns true if only the first lane of \p Def is used.
bool onlyFirstLaneUsed(const VPValue *Def);
@@ -159,6 +181,19 @@ bool isUsedByLoadStoreAddress(const VPValue *V);
/// inserted for predicated reductions or tail folding.
VPInstruction *findComputeReductionResult(VPReductionPHIRecipe *PhiR);
+/// Return an add/fadd or accumulator-minus-extended-op sub/fsub IR reduction
+/// chain that can be represented using partial reductions. This is a
+/// conservative subset of VPlan partial-reduction formation used when probing
+/// whether partial reductions can affect VF limits.
+LLVM_ABI std::optional<SmallVector<PartialReductionChainLink>>
+getPartialReductionChain(PHINode *Phi, const RecurrenceDescriptor &RdxDesc,
+ const Loop *L);
+
+/// Return the target cost for the given partial-reduction chain link at \p VF.
+LLVM_ABI InstructionCost getPartialReductionCost(
+ const PartialReductionChainLink &Link, const TargetTransformInfo &TTI,
+ ElementCount VF, TTI::TargetCostKind CostKind);
+
/// Collect the header mask with the pattern:
/// (ICMP_ULE, WideCanonicalIV, backedge-taken-count)
/// Note: If alias masking is enabled this will find:
diff --git a/llvm/test/CodeGen/X86/partial-reduce-avx512-vex-split.ll b/llvm/test/CodeGen/X86/partial-reduce-avx512-vex-split.ll
new file mode 100644
index 0000000000000..ca32194f50bc1
--- /dev/null
+++ b/llvm/test/CodeGen/X86/partial-reduce-avx512-vex-split.ll
@@ -0,0 +1,89 @@
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+avx512f,+avx512bw,+avx512vl,+avxvnni,-avx512vnni < %s | FileCheck %s --check-prefix=AVXVNNI
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+avx512f,+avx512bw,+avx512vl,+avxvnniint8,-avx10.2 < %s | FileCheck %s --check-prefix=AVXVNNIINT8
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+avx512f,+avx512bw,+avx512vl,+avxvnniint16,-avx10.2 < %s | FileCheck %s --check-prefix=AVXVNNIINT16
+
+; Targets with AVX512 registers but only 256-bit VEX/EVEX dot-product
+; instructions should split 512-bit partial reductions to two 256-bit dot
+; products instead of expanding them as scalarized zmm arithmetic.
+
+define <16 x i32> @partial_reduce_sumla_i8_v16i32(<16 x i32> %acc, <64 x i8> %a, <64 x i8> %b) {
+; AVXVNNI-LABEL: partial_reduce_sumla_i8_v16i32:
+; AVXVNNI: {vex} vpdpbusd
+; AVXVNNI: {vex} vpdpbusd
+; AVXVNNI-NOT: vpmaddwd
+; AVXVNNI: retq
+ %a.zext = zext <64 x i8> %a to <64 x i32>
+ %b.sext = sext <64 x i8> %b to <64 x i32>
+ %mul = mul nsw <64 x i32> %a.zext, %b.sext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> %acc, <64 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+define <16 x i32> @partial_reduce_smla_i16_v16i32(<16 x i32> %acc, <32 x i16> %a, <32 x i16> %b) {
+; AVXVNNI-LABEL: partial_reduce_smla_i16_v16i32:
+; AVXVNNI: {vex} vpdpwssd
+; AVXVNNI: vpmaddwd
+; AVXVNNI: vpaddd
+; AVXVNNI-NOT: vpmulld
+; AVXVNNI: retq
+ %a.sext = sext <32 x i16> %a to <32 x i32>
+ %b.sext = sext <32 x i16> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.sext, %b.sext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> %acc, <32 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+define <16 x i32> @partial_reduce_smla_i8_v16i32(<16 x i32> %acc, <64 x i8> %a, <64 x i8> %b) {
+; AVXVNNIINT8-LABEL: partial_reduce_smla_i8_v16i32:
+; AVXVNNIINT8: vpdpbssd
+; AVXVNNIINT8: vpdpbssd
+; AVXVNNIINT8-NOT: vpmulld
+; AVXVNNIINT8: retq
+ %a.sext = sext <64 x i8> %a to <64 x i32>
+ %b.sext = sext <64 x i8> %b to <64 x i32>
+ %mul = mul nsw <64 x i32> %a.sext, %b.sext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> %acc, <64 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+define <16 x i32> @partial_reduce_umla_i8_v16i32(<16 x i32> %acc, <64 x i8> %a, <64 x i8> %b) {
+; AVXVNNIINT8-LABEL: partial_reduce_umla_i8_v16i32:
+; AVXVNNIINT8: vpdpbuud
+; AVXVNNIINT8: vpdpbuud
+; AVXVNNIINT8-NOT: vpmaddwd
+; AVXVNNIINT8: retq
+ %a.zext = zext <64 x i8> %a to <64 x i32>
+ %b.zext = zext <64 x i8> %b to <64 x i32>
+ %mul = mul nsw <64 x i32> %a.zext, %b.zext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> %acc, <64 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+define <16 x i32> @partial_reduce_sumla_i16_v16i32(<16 x i32> %acc, <32 x i16> %a, <32 x i16> %b) {
+; AVXVNNIINT16-LABEL: partial_reduce_sumla_i16_v16i32:
+; AVXVNNIINT16: vpdpwsud
+; AVXVNNIINT16: vpdpwsud
+; AVXVNNIINT16-NOT: vpmulld
+; AVXVNNIINT16: retq
+ %a.sext = sext <32 x i16> %a to <32 x i32>
+ %b.zext = zext <32 x i16> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.sext, %b.zext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> %acc, <32 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+define <16 x i32> @partial_reduce_umla_i16_v16i32(<16 x i32> %acc, <32 x i16> %a, <32 x i16> %b) {
+; AVXVNNIINT16-LABEL: partial_reduce_umla_i16_v16i32:
+; AVXVNNIINT16: vpdpwuud
+; AVXVNNIINT16: vpdpwuud
+; AVXVNNIINT16-NOT: vpmulld
+; AVXVNNIINT16: retq
+ %a.zext = zext <32 x i16> %a to <32 x i32>
+ %b.zext = zext <32 x i16> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.zext, %b.zext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> %acc, <32 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+declare <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32>, <64 x i32>)
+declare <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32>, <32 x i32>)
diff --git a/llvm/test/CodeGen/X86/partial-reduce-dot-product.ll b/llvm/test/CodeGen/X86/partial-reduce-dot-product.ll
new file mode 100644
index 0000000000000..7871e95335ef0
--- /dev/null
+++ b/llvm/test/CodeGen/X86/partial-reduce-dot-product.ll
@@ -0,0 +1,1033 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=znver4 < %s | FileCheck %s --check-prefixes=AVX512VNNI
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=alderlake < %s | FileCheck %s --check-prefixes=AVXVNNI
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=znver4 -mattr=+avxvnniint8,+avxvnniint16 < %s | FileCheck %s --check-prefixes=AVXVNNIINT8INT16
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+avx10.2-512 < %s | FileCheck %s --check-prefixes=AVX10
+
+; Test that PARTIAL_REDUCE_SUMLA (sext × zext) lowers to vpdpbusd.
+define <4 x i32> @partial_reduce_sumla_v4i32(<4 x i32> %acc, <16 x i8> %a, <16 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_sumla_v4i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpdpbusd %xmm2, %xmm1, %xmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_sumla_v4i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: {vex} vpdpbusd %xmm2, %xmm1, %xmm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_sumla_v4i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbusd %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_sumla_v4i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbusd %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.zext = zext <16 x i8> %a to <16 x i32>
+ %b.sext = sext <16 x i8> %b to <16 x i32>
+ %mul = mul nsw <16 x i32> %a.zext, %b.sext
+ %res = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> %acc, <16 x i32> %mul)
+ ret <4 x i32> %res
+}
+
+; Test 512-bit (zmm) partial reduction.
+define <16 x i32> @partial_reduce_sumla_v16i32(<16 x i32> %acc, <64 x i8> %a, <64 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_sumla_v16i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpdpbusd %zmm2, %zmm1, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_sumla_v16i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: {vex} vpdpbusd %ymm4, %ymm2, %ymm0
+; AVXVNNI-NEXT: {vex} vpdpbusd %ymm5, %ymm3, %ymm1
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_sumla_v16i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbusd %zmm2, %zmm1, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_sumla_v16i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbusd %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.zext = zext <64 x i8> %a to <64 x i32>
+ %b.sext = sext <64 x i8> %b to <64 x i32>
+ %mul = mul nsw <64 x i32> %a.zext, %b.sext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> %acc, <64 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+; Test 256-bit (ymm) partial reduction.
+define <8 x i32> @partial_reduce_sumla_v8i32(<8 x i32> %acc, <32 x i8> %a, <32 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_sumla_v8i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpdpbusd %ymm2, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_sumla_v8i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: {vex} vpdpbusd %ymm2, %ymm1, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_sumla_v8i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbusd %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_sumla_v8i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbusd %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.zext = zext <32 x i8> %a to <32 x i32>
+ %b.sext = sext <32 x i8> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.zext, %b.sext
+ %res = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v32i32(<8 x i32> %acc, <32 x i32> %mul)
+ ret <8 x i32> %res
+}
+
+declare <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32>, <16 x i32>)
+declare <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v32i32(<8 x i32>, <32 x i32>)
+declare <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32>, <64 x i32>)
+
+; i16 x i16 -> i32 partial reduction tests (vpdpwssd)
+
+; Test 128-bit (xmm) i16 partial reduction: 8 x i16 -> 4 x i32.
+define <4 x i32> @partial_reduce_smla_i16_v4i32(<4 x i32> %acc, <8 x i16> %a, <8 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_smla_i16_v4i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpdpwssd %xmm2, %xmm1, %xmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_smla_i16_v4i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: {vex} vpdpwssd %xmm2, %xmm1, %xmm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_smla_i16_v4i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwssd %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_smla_i16_v4i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwssd %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.sext = sext <8 x i16> %a to <8 x i32>
+ %b.sext = sext <8 x i16> %b to <8 x i32>
+ %mul = mul nsw <8 x i32> %a.sext, %b.sext
+ %res = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v8i32(<4 x i32> %acc, <8 x i32> %mul)
+ ret <4 x i32> %res
+}
+
+; Test 256-bit (ymm) i16 partial reduction: 16 x i16 -> 8 x i32.
+define <8 x i32> @partial_reduce_smla_i16_v8i32(<8 x i32> %acc, <16 x i16> %a, <16 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_smla_i16_v8i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpdpwssd %ymm2, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_smla_i16_v8i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: {vex} vpdpwssd %ymm2, %ymm1, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_smla_i16_v8i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwssd %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_smla_i16_v8i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwssd %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.sext = sext <16 x i16> %a to <16 x i32>
+ %b.sext = sext <16 x i16> %b to <16 x i32>
+ %mul = mul nsw <16 x i32> %a.sext, %b.sext
+ %res = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v16i32(<8 x i32> %acc, <16 x i32> %mul)
+ ret <8 x i32> %res
+}
+
+; Test 512-bit (zmm) i16 partial reduction: 32 x i16 -> 16 x i32.
+define <16 x i32> @partial_reduce_smla_i16_v16i32(<16 x i32> %acc, <32 x i16> %a, <32 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_smla_i16_v16i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpdpwssd %zmm2, %zmm1, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_smla_i16_v16i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: {vex} vpdpwssd %ymm4, %ymm2, %ymm0
+; AVXVNNI-NEXT: {vex} vpdpwssd %ymm5, %ymm3, %ymm1
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_smla_i16_v16i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwssd %zmm2, %zmm1, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_smla_i16_v16i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwssd %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.sext = sext <32 x i16> %a to <32 x i32>
+ %b.sext = sext <32 x i16> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.sext, %b.sext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> %acc, <32 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+; bf16 x bf16 -> f32 partial reduction tests (vdpbf16ps)
+
+; Test 128-bit (xmm) bf16 partial reduction: 8 x bf16 -> 4 x f32.
+define <4 x float> @partial_reduce_fmla_bf16_v4f32(<4 x float> %acc, <8 x bfloat> %a, <8 x bfloat> %b) {
+; AVX512VNNI-LABEL: partial_reduce_fmla_bf16_v4f32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vdpbf16ps %xmm2, %xmm1, %xmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_fmla_bf16_v4f32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm1 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm2, %ymm2
+; AVXVNNI-NEXT: vmulps %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vaddps %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT: vextractf128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vaddps %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT: vzeroupper
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_fmla_bf16_v4f32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vdpbf16ps %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_fmla_bf16_v4f32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vdpbf16ps %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.ext = fpext <8 x bfloat> %a to <8 x float>
+ %b.ext = fpext <8 x bfloat> %b to <8 x float>
+ %mul = fmul <8 x float> %a.ext, %b.ext
+ %res = call <4 x float> @llvm.vector.partial.reduce.fadd.v4f32.v8f32(<4 x float> %acc, <8 x float> %mul)
+ ret <4 x float> %res
+}
+
+; Test 256-bit (ymm) bf16 partial reduction: 16 x bf16 -> 8 x f32.
+define <8 x float> @partial_reduce_fmla_bf16_v8f32(<8 x float> %acc, <16 x bfloat> %a, <16 x bfloat> %b) {
+; AVX512VNNI-LABEL: partial_reduce_fmla_bf16_v8f32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vdpbf16ps %ymm2, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_fmla_bf16_v8f32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm3 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm3, %ymm3
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm1 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm4, %ymm4
+; AVXVNNI-NEXT: vmulps %ymm4, %ymm3, %ymm3
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm2, %ymm2
+; AVXVNNI-NEXT: vmulps %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vaddps %ymm3, %ymm0, %ymm0
+; AVXVNNI-NEXT: vaddps %ymm1, %ymm0, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_fmla_bf16_v8f32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vdpbf16ps %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_fmla_bf16_v8f32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vdpbf16ps %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.ext = fpext <16 x bfloat> %a to <16 x float>
+ %b.ext = fpext <16 x bfloat> %b to <16 x float>
+ %mul = fmul <16 x float> %a.ext, %b.ext
+ %res = call <8 x float> @llvm.vector.partial.reduce.fadd.v8f32.v16f32(<8 x float> %acc, <16 x float> %mul)
+ ret <8 x float> %res
+}
+
+; Test 512-bit (zmm) bf16 partial reduction: 32 x bf16 -> 16 x f32.
+define <16 x float> @partial_reduce_fmla_bf16_v16f32(<16 x float> %acc, <32 x bfloat> %a, <32 x bfloat> %b) {
+; AVX512VNNI-LABEL: partial_reduce_fmla_bf16_v16f32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vdpbf16ps %zmm2, %zmm1, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_fmla_bf16_v16f32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm6 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm6, %ymm6
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm2, %ymm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm7 = xmm3[0],zero,xmm3[1],zero,xmm3[2],zero,xmm3[3],zero,xmm3[4],zero,xmm3[5],zero,xmm3[6],zero,xmm3[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm7, %ymm7
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm3
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm3 = xmm3[0],zero,xmm3[1],zero,xmm3[2],zero,xmm3[3],zero,xmm3[4],zero,xmm3[5],zero,xmm3[6],zero,xmm3[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm8 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm8, %ymm8
+; AVXVNNI-NEXT: vmulps %ymm6, %ymm8, %ymm6
+; AVXVNNI-NEXT: vextracti128 $1, %ymm4, %xmm4
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm4, %ymm4
+; AVXVNNI-NEXT: vmulps %ymm4, %ymm2, %ymm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm5[0],zero,xmm5[1],zero,xmm5[2],zero,xmm5[3],zero,xmm5[4],zero,xmm5[5],zero,xmm5[6],zero,xmm5[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm4, %ymm4
+; AVXVNNI-NEXT: vmulps %ymm4, %ymm7, %ymm4
+; AVXVNNI-NEXT: vextracti128 $1, %ymm5, %xmm5
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm5 = xmm5[0],zero,xmm5[1],zero,xmm5[2],zero,xmm5[3],zero,xmm5[4],zero,xmm5[5],zero,xmm5[6],zero,xmm5[7],zero
+; AVXVNNI-NEXT: vpslld $16, %ymm5, %ymm5
+; AVXVNNI-NEXT: vmulps %ymm5, %ymm3, %ymm3
+; AVXVNNI-NEXT: vaddps %ymm6, %ymm0, %ymm0
+; AVXVNNI-NEXT: vaddps %ymm2, %ymm0, %ymm0
+; AVXVNNI-NEXT: vaddps %ymm4, %ymm1, %ymm1
+; AVXVNNI-NEXT: vaddps %ymm3, %ymm1, %ymm1
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_fmla_bf16_v16f32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vdpbf16ps %zmm2, %zmm1, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_fmla_bf16_v16f32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vdpbf16ps %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.ext = fpext <32 x bfloat> %a to <32 x float>
+ %b.ext = fpext <32 x bfloat> %b to <32 x float>
+ %mul = fmul <32 x float> %a.ext, %b.ext
+ %res = call <16 x float> @llvm.vector.partial.reduce.fadd.v16f32.v32f32(<16 x float> %acc, <32 x float> %mul)
+ ret <16 x float> %res
+}
+
+; VNNI-INT8 i8 x i8 -> i32 SMLA (sext x sext) tests: VPDPBSSD
+
+; Test 128-bit (xmm) SMLA i8.
+define <4 x i32> @partial_reduce_smla_i8_v4i32(<4 x i32> %acc, <16 x i8> %a, <16 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_smla_i8_v4i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovsxbd %xmm1, %zmm1
+; AVX512VNNI-NEXT: vpmovsxbd %xmm2, %zmm2
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT: vextracti32x4 $3, %zmm1, %xmm2
+; AVX512VNNI-NEXT: vextracti32x4 $2, %zmm1, %xmm3
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT: vpaddd %xmm3, %xmm1, %xmm1
+; AVX512VNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VNNI-NEXT: vzeroupper
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_smla_i8_v4i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovsxbd %xmm1, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm1 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm1, %ymm1
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm4
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm2 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm2
+; AVXVNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %xmm3, %xmm0, %xmm0
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm2
+; AVXVNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT: vpaddd %xmm0, %xmm2, %xmm0
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT: vzeroupper
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_smla_i8_v4i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbssd %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_smla_i8_v4i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbssd %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.sext = sext <16 x i8> %a to <16 x i32>
+ %b.sext = sext <16 x i8> %b to <16 x i32>
+ %mul = mul nsw <16 x i32> %a.sext, %b.sext
+ %res = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> %acc, <16 x i32> %mul)
+ ret <4 x i32> %res
+}
+
+; Test 256-bit (ymm) SMLA i8.
+define <8 x i32> @partial_reduce_smla_i8_v8i32(<8 x i32> %acc, <32 x i8> %a, <32 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_smla_i8_v8i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovsxbd %xmm1, %zmm3
+; AVX512VNNI-NEXT: vpmovsxbd %xmm2, %zmm4
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVX512VNNI-NEXT: vpmovsxbd %xmm1, %zmm1
+; AVX512VNNI-NEXT: vpmovsxbd %xmm2, %zmm2
+; AVX512VNNI-NEXT: vpmulld %zmm4, %zmm3, %zmm3
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm3, %ymm2
+; AVX512VNNI-NEXT: vpaddd %ymm0, %ymm2, %ymm0
+; AVX512VNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vpaddd %ymm0, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_smla_i8_v8i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovsxbd %xmm1, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm4 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm4, %ymm4
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpmovsxbd %xmm1, %ymm5
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm1 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm1, %ymm1
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm6
+; AVXVNNI-NEXT: vpmulld %ymm6, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm6 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm6, %ymm6
+; AVXVNNI-NEXT: vpmulld %ymm6, %ymm4, %ymm4
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm6
+; AVXVNNI-NEXT: vpmulld %ymm6, %ymm5, %ymm5
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm2 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm2
+; AVXVNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm4, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm5, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_smla_i8_v8i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbssd %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_smla_i8_v8i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbssd %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.sext = sext <32 x i8> %a to <32 x i32>
+ %b.sext = sext <32 x i8> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.sext, %b.sext
+ %res = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v32i32(<8 x i32> %acc, <32 x i32> %mul)
+ ret <8 x i32> %res
+}
+
+; Test 512-bit (zmm) SMLA i8.
+; VEX (AVXVNNIINT8): splits to 2x256-bit ymm.
+; EVEX (AVX10.2): native 512-bit zmm.
+define <16 x i32> @partial_reduce_smla_i8_v16i32(<16 x i32> %acc, <64 x i8> %a, <64 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_smla_i8_v16i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovsxbd %xmm1, %zmm3
+; AVX512VNNI-NEXT: vpmovsxbd %xmm2, %zmm6
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm4
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm2, %xmm7
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm2, %ymm2
+; AVX512VNNI-NEXT: vpmovsxbd %xmm4, %zmm4
+; AVX512VNNI-NEXT: vpmovsxbd %xmm1, %zmm5
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vpmovsxbd %xmm1, %zmm1
+; AVX512VNNI-NEXT: vpmulld %zmm6, %zmm3, %zmm3
+; AVX512VNNI-NEXT: vpmovsxbd %xmm7, %zmm6
+; AVX512VNNI-NEXT: vpmovsxbd %xmm2, %zmm7
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVX512VNNI-NEXT: vpmovsxbd %xmm2, %zmm2
+; AVX512VNNI-NEXT: vpaddd %zmm3, %zmm0, %zmm0
+; AVX512VNNI-NEXT: vpmulld %zmm6, %zmm4, %zmm4
+; AVX512VNNI-NEXT: vpmulld %zmm7, %zmm5, %zmm5
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %zmm4, %zmm0, %zmm0
+; AVX512VNNI-NEXT: vpaddd %zmm5, %zmm0, %zmm0
+; AVX512VNNI-NEXT: vpaddd %zmm1, %zmm0, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_smla_i8_v16i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm6
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm7 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm7, %ymm7
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm8
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm2 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm2, %ymm2
+; AVXVNNI-NEXT: vpmovsxbd %xmm3, %ymm9
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm10 = xmm3[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm10, %ymm10
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm3
+; AVXVNNI-NEXT: vpmovsxbd %xmm3, %ymm11
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm3 = xmm3[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm4, %ymm12
+; AVXVNNI-NEXT: vpmulld %ymm12, %ymm6, %ymm6
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm12 = xmm4[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm12, %ymm12
+; AVXVNNI-NEXT: vpmulld %ymm12, %ymm7, %ymm7
+; AVXVNNI-NEXT: vextracti128 $1, %ymm4, %xmm4
+; AVXVNNI-NEXT: vpmovsxbd %xmm4, %ymm12
+; AVXVNNI-NEXT: vpmulld %ymm12, %ymm8, %ymm8
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm4 = xmm4[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm4, %ymm4
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm2, %ymm2
+; AVXVNNI-NEXT: vpmovsxbd %xmm3, %ymm3
+; AVXVNNI-NEXT: vpmovsxbd %xmm5, %ymm4
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm9, %ymm4
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm9 = xmm5[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm9, %ymm9
+; AVXVNNI-NEXT: vpmulld %ymm9, %ymm10, %ymm9
+; AVXVNNI-NEXT: vextracti128 $1, %ymm5, %xmm5
+; AVXVNNI-NEXT: vpmovsxbd %xmm5, %ymm10
+; AVXVNNI-NEXT: vpmulld %ymm10, %ymm11, %ymm10
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm5 = xmm5[2,3,2,3]
+; AVXVNNI-NEXT: vpmovsxbd %xmm5, %ymm5
+; AVXVNNI-NEXT: vpmulld %ymm5, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpaddd %ymm6, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm7, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm0, %ymm8, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm2, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm4, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm1, %ymm9, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm1, %ymm10, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm1, %ymm1
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_smla_i8_v16i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm0, %ymm5
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm2, %ymm3
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm1, %ymm4
+; AVXVNNIINT8INT16-NEXT: vpdpbssd %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: vpdpbssd %ymm3, %ymm4, %ymm5
+; AVXVNNIINT8INT16-NEXT: vinsertf64x4 $1, %ymm5, %zmm0, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_smla_i8_v16i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbssd %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.sext = sext <64 x i8> %a to <64 x i32>
+ %b.sext = sext <64 x i8> %b to <64 x i32>
+ %mul = mul nsw <64 x i32> %a.sext, %b.sext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> %acc, <64 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+; VNNI-INT8 i8 x i8 -> i32 UMLA (zext x zext) tests: VPDPBUUD
+
+; Test 128-bit (xmm) UMLA i8.
+define <4 x i32> @partial_reduce_umla_i8_v4i32(<4 x i32> %acc, <16 x i8> %a, <16 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_umla_i8_v4i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm1 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero,xmm1[8],zero,zero,zero,xmm1[9],zero,zero,zero,xmm1[10],zero,zero,zero,xmm1[11],zero,zero,zero,xmm1[12],zero,zero,zero,xmm1[13],zero,zero,zero,xmm1[14],zero,zero,zero,xmm1[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm2 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero,xmm2[8],zero,zero,zero,xmm2[9],zero,zero,zero,xmm2[10],zero,zero,zero,xmm2[11],zero,zero,zero,xmm2[12],zero,zero,zero,xmm2[13],zero,zero,zero,xmm2[14],zero,zero,zero,xmm2[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmaddwd %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT: vextracti32x4 $3, %zmm1, %xmm2
+; AVX512VNNI-NEXT: vextracti32x4 $2, %zmm1, %xmm3
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vpaddd %xmm0, %xmm2, %xmm0
+; AVX512VNNI-NEXT: vpaddd %xmm3, %xmm1, %xmm1
+; AVX512VNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VNNI-NEXT: vzeroupper
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_umla_i8_v4i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm3 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm1 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm1 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm4 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm4, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm2 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm2 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %xmm3, %xmm0, %xmm0
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm2
+; AVXVNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT: vpaddd %xmm0, %xmm2, %xmm0
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT: vzeroupper
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_umla_i8_v4i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbuud %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_umla_i8_v4i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbuud %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.zext = zext <16 x i8> %a to <16 x i32>
+ %b.zext = zext <16 x i8> %b to <16 x i32>
+ %mul = mul nsw <16 x i32> %a.zext, %b.zext
+ %res = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32(<4 x i32> %acc, <16 x i32> %mul)
+ ret <4 x i32> %res
+}
+
+; Test 256-bit (ymm) UMLA i8.
+define <8 x i32> @partial_reduce_umla_i8_v8i32(<8 x i32> %acc, <32 x i8> %a, <32 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_umla_i8_v8i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm3 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero,xmm1[8],zero,zero,zero,xmm1[9],zero,zero,zero,xmm1[10],zero,zero,zero,xmm1[11],zero,zero,zero,xmm1[12],zero,zero,zero,xmm1[13],zero,zero,zero,xmm1[14],zero,zero,zero,xmm1[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm4 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero,xmm2[8],zero,zero,zero,xmm2[9],zero,zero,zero,xmm2[10],zero,zero,zero,xmm2[11],zero,zero,zero,xmm2[12],zero,zero,zero,xmm2[13],zero,zero,zero,xmm2[14],zero,zero,zero,xmm2[15],zero,zero,zero
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm1 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero,xmm1[8],zero,zero,zero,xmm1[9],zero,zero,zero,xmm1[10],zero,zero,zero,xmm1[11],zero,zero,zero,xmm1[12],zero,zero,zero,xmm1[13],zero,zero,zero,xmm1[14],zero,zero,zero,xmm1[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm2 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero,xmm2[8],zero,zero,zero,xmm2[9],zero,zero,zero,xmm2[10],zero,zero,zero,xmm2[11],zero,zero,zero,xmm2[12],zero,zero,zero,xmm2[13],zero,zero,zero,xmm2[14],zero,zero,zero,xmm2[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmaddwd %zmm4, %zmm3, %zmm3
+; AVX512VNNI-NEXT: vpmaddwd %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm3, %ymm2
+; AVX512VNNI-NEXT: vpaddd %ymm0, %ymm2, %ymm0
+; AVX512VNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vpaddd %ymm0, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_umla_i8_v8i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm4 = xmm3[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm4 = xmm4[0],zero,zero,zero,xmm4[1],zero,zero,zero,xmm4[2],zero,zero,zero,xmm4[3],zero,zero,zero,xmm4[4],zero,zero,zero,xmm4[5],zero,zero,zero,xmm4[6],zero,zero,zero,xmm4[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm3 = xmm3[0],zero,zero,zero,xmm3[1],zero,zero,zero,xmm3[2],zero,zero,zero,xmm3[3],zero,zero,zero,xmm3[4],zero,zero,zero,xmm3[5],zero,zero,zero,xmm3[6],zero,zero,zero,xmm3[7],zero,zero,zero
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm5 = xmm1[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm5 = xmm5[0],zero,zero,zero,xmm5[1],zero,zero,zero,xmm5[2],zero,zero,zero,xmm5[3],zero,zero,zero,xmm5[4],zero,zero,zero,xmm5[5],zero,zero,zero,xmm5[6],zero,zero,zero,xmm5[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm1 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm6
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm7 = xmm6[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm7 = xmm7[0],zero,zero,zero,xmm7[1],zero,zero,zero,xmm7[2],zero,zero,zero,xmm7[3],zero,zero,zero,xmm7[4],zero,zero,zero,xmm7[5],zero,zero,zero,xmm7[6],zero,zero,zero,xmm7[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm7, %ymm4, %ymm4
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm6 = xmm6[0],zero,zero,zero,xmm6[1],zero,zero,zero,xmm6[2],zero,zero,zero,xmm6[3],zero,zero,zero,xmm6[4],zero,zero,zero,xmm6[5],zero,zero,zero,xmm6[6],zero,zero,zero,xmm6[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm6, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm6 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm6 = xmm6[0],zero,zero,zero,xmm6[1],zero,zero,zero,xmm6[2],zero,zero,zero,xmm6[3],zero,zero,zero,xmm6[4],zero,zero,zero,xmm6[5],zero,zero,zero,xmm6[6],zero,zero,zero,xmm6[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm6, %ymm5, %ymm5
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm2 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero
+; AVXVNNI-NEXT: {vex} vpdpwssd %ymm2, %ymm1, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm5, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm4, %ymm0, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_umla_i8_v8i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpbuud %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_umla_i8_v8i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbuud %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.zext = zext <32 x i8> %a to <32 x i32>
+ %b.zext = zext <32 x i8> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.zext, %b.zext
+ %res = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v32i32(<8 x i32> %acc, <32 x i32> %mul)
+ ret <8 x i32> %res
+}
+
+; Test 512-bit (zmm) UMLA i8.
+define <16 x i32> @partial_reduce_umla_i8_v16i32(<16 x i32> %acc, <64 x i8> %a, <64 x i8> %b) {
+; AVX512VNNI-LABEL: partial_reduce_umla_i8_v16i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm5
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm3
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm2, %xmm8
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm2, %ymm6
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm1 = xmm1[0],zero,zero,zero,xmm1[1],zero,zero,zero,xmm1[2],zero,zero,zero,xmm1[3],zero,zero,zero,xmm1[4],zero,zero,zero,xmm1[5],zero,zero,zero,xmm1[6],zero,zero,zero,xmm1[7],zero,zero,zero,xmm1[8],zero,zero,zero,xmm1[9],zero,zero,zero,xmm1[10],zero,zero,zero,xmm1[11],zero,zero,zero,xmm1[12],zero,zero,zero,xmm1[13],zero,zero,zero,xmm1[14],zero,zero,zero,xmm1[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm2 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero,xmm2[8],zero,zero,zero,xmm2[9],zero,zero,zero,xmm2[10],zero,zero,zero,xmm2[11],zero,zero,zero,xmm2[12],zero,zero,zero,xmm2[13],zero,zero,zero,xmm2[14],zero,zero,zero,xmm2[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm5 = xmm5[0],zero,zero,zero,xmm5[1],zero,zero,zero,xmm5[2],zero,zero,zero,xmm5[3],zero,zero,zero,xmm5[4],zero,zero,zero,xmm5[5],zero,zero,zero,xmm5[6],zero,zero,zero,xmm5[7],zero,zero,zero,xmm5[8],zero,zero,zero,xmm5[9],zero,zero,zero,xmm5[10],zero,zero,zero,xmm5[11],zero,zero,zero,xmm5[12],zero,zero,zero,xmm5[13],zero,zero,zero,xmm5[14],zero,zero,zero,xmm5[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm8 = xmm8[0],zero,zero,zero,xmm8[1],zero,zero,zero,xmm8[2],zero,zero,zero,xmm8[3],zero,zero,zero,xmm8[4],zero,zero,zero,xmm8[5],zero,zero,zero,xmm8[6],zero,zero,zero,xmm8[7],zero,zero,zero,xmm8[8],zero,zero,zero,xmm8[9],zero,zero,zero,xmm8[10],zero,zero,zero,xmm8[11],zero,zero,zero,xmm8[12],zero,zero,zero,xmm8[13],zero,zero,zero,xmm8[14],zero,zero,zero,xmm8[15],zero,zero,zero
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm3, %xmm4
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm6, %xmm7
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm3 = xmm3[0],zero,zero,zero,xmm3[1],zero,zero,zero,xmm3[2],zero,zero,zero,xmm3[3],zero,zero,zero,xmm3[4],zero,zero,zero,xmm3[5],zero,zero,zero,xmm3[6],zero,zero,zero,xmm3[7],zero,zero,zero,xmm3[8],zero,zero,zero,xmm3[9],zero,zero,zero,xmm3[10],zero,zero,zero,xmm3[11],zero,zero,zero,xmm3[12],zero,zero,zero,xmm3[13],zero,zero,zero,xmm3[14],zero,zero,zero,xmm3[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm6 = xmm6[0],zero,zero,zero,xmm6[1],zero,zero,zero,xmm6[2],zero,zero,zero,xmm6[3],zero,zero,zero,xmm6[4],zero,zero,zero,xmm6[5],zero,zero,zero,xmm6[6],zero,zero,zero,xmm6[7],zero,zero,zero,xmm6[8],zero,zero,zero,xmm6[9],zero,zero,zero,xmm6[10],zero,zero,zero,xmm6[11],zero,zero,zero,xmm6[12],zero,zero,zero,xmm6[13],zero,zero,zero,xmm6[14],zero,zero,zero,xmm6[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm4 = xmm4[0],zero,zero,zero,xmm4[1],zero,zero,zero,xmm4[2],zero,zero,zero,xmm4[3],zero,zero,zero,xmm4[4],zero,zero,zero,xmm4[5],zero,zero,zero,xmm4[6],zero,zero,zero,xmm4[7],zero,zero,zero,xmm4[8],zero,zero,zero,xmm4[9],zero,zero,zero,xmm4[10],zero,zero,zero,xmm4[11],zero,zero,zero,xmm4[12],zero,zero,zero,xmm4[13],zero,zero,zero,xmm4[14],zero,zero,zero,xmm4[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpmovzxbd {{.*#+}} zmm7 = xmm7[0],zero,zero,zero,xmm7[1],zero,zero,zero,xmm7[2],zero,zero,zero,xmm7[3],zero,zero,zero,xmm7[4],zero,zero,zero,xmm7[5],zero,zero,zero,xmm7[6],zero,zero,zero,xmm7[7],zero,zero,zero,xmm7[8],zero,zero,zero,xmm7[9],zero,zero,zero,xmm7[10],zero,zero,zero,xmm7[11],zero,zero,zero,xmm7[12],zero,zero,zero,xmm7[13],zero,zero,zero,xmm7[14],zero,zero,zero,xmm7[15],zero,zero,zero
+; AVX512VNNI-NEXT: vpdpwssd %zmm2, %zmm1, %zmm0
+; AVX512VNNI-NEXT: vpdpwssd %zmm8, %zmm5, %zmm0
+; AVX512VNNI-NEXT: vpdpwssd %zmm6, %zmm3, %zmm0
+; AVX512VNNI-NEXT: vpdpwssd %zmm7, %zmm4, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_umla_i8_v16i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm6
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm7 = xmm6[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm7 = xmm7[0],zero,zero,zero,xmm7[1],zero,zero,zero,xmm7[2],zero,zero,zero,xmm7[3],zero,zero,zero,xmm7[4],zero,zero,zero,xmm7[5],zero,zero,zero,xmm7[6],zero,zero,zero,xmm7[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm6 = xmm6[0],zero,zero,zero,xmm6[1],zero,zero,zero,xmm6[2],zero,zero,zero,xmm6[3],zero,zero,zero,xmm6[4],zero,zero,zero,xmm6[5],zero,zero,zero,xmm6[6],zero,zero,zero,xmm6[7],zero,zero,zero
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm8 = xmm3[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm8 = xmm8[0],zero,zero,zero,xmm8[1],zero,zero,zero,xmm8[2],zero,zero,zero,xmm8[3],zero,zero,zero,xmm8[4],zero,zero,zero,xmm8[5],zero,zero,zero,xmm8[6],zero,zero,zero,xmm8[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm3 = xmm3[0],zero,zero,zero,xmm3[1],zero,zero,zero,xmm3[2],zero,zero,zero,xmm3[3],zero,zero,zero,xmm3[4],zero,zero,zero,xmm3[5],zero,zero,zero,xmm3[6],zero,zero,zero,xmm3[7],zero,zero,zero
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm9
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm10 = xmm9[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm10 = xmm10[0],zero,zero,zero,xmm10[1],zero,zero,zero,xmm10[2],zero,zero,zero,xmm10[3],zero,zero,zero,xmm10[4],zero,zero,zero,xmm10[5],zero,zero,zero,xmm10[6],zero,zero,zero,xmm10[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm9 = xmm9[0],zero,zero,zero,xmm9[1],zero,zero,zero,xmm9[2],zero,zero,zero,xmm9[3],zero,zero,zero,xmm9[4],zero,zero,zero,xmm9[5],zero,zero,zero,xmm9[6],zero,zero,zero,xmm9[7],zero,zero,zero
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm11 = xmm2[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm11 = xmm11[0],zero,zero,zero,xmm11[1],zero,zero,zero,xmm11[2],zero,zero,zero,xmm11[3],zero,zero,zero,xmm11[4],zero,zero,zero,xmm11[5],zero,zero,zero,xmm11[6],zero,zero,zero,xmm11[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm2 = xmm2[0],zero,zero,zero,xmm2[1],zero,zero,zero,xmm2[2],zero,zero,zero,xmm2[3],zero,zero,zero,xmm2[4],zero,zero,zero,xmm2[5],zero,zero,zero,xmm2[6],zero,zero,zero,xmm2[7],zero,zero,zero
+; AVXVNNI-NEXT: vextracti128 $1, %ymm5, %xmm12
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm13 = xmm12[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm13 = xmm13[0],zero,zero,zero,xmm13[1],zero,zero,zero,xmm13[2],zero,zero,zero,xmm13[3],zero,zero,zero,xmm13[4],zero,zero,zero,xmm13[5],zero,zero,zero,xmm13[6],zero,zero,zero,xmm13[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm7, %ymm13, %ymm7
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm12 = xmm12[0],zero,zero,zero,xmm12[1],zero,zero,zero,xmm12[2],zero,zero,zero,xmm12[3],zero,zero,zero,xmm12[4],zero,zero,zero,xmm12[5],zero,zero,zero,xmm12[6],zero,zero,zero,xmm12[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm6, %ymm12, %ymm6
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm12 = xmm5[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm12 = xmm12[0],zero,zero,zero,xmm12[1],zero,zero,zero,xmm12[2],zero,zero,zero,xmm12[3],zero,zero,zero,xmm12[4],zero,zero,zero,xmm12[5],zero,zero,zero,xmm12[6],zero,zero,zero,xmm12[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm12, %ymm8, %ymm8
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm5 = xmm5[0],zero,zero,zero,xmm5[1],zero,zero,zero,xmm5[2],zero,zero,zero,xmm5[3],zero,zero,zero,xmm5[4],zero,zero,zero,xmm5[5],zero,zero,zero,xmm5[6],zero,zero,zero,xmm5[7],zero,zero,zero
+; AVXVNNI-NEXT: {vex} vpdpwssd %ymm5, %ymm3, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm1, %ymm8, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm6, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm7, %ymm1, %ymm1
+; AVXVNNI-NEXT: vextracti128 $1, %ymm4, %xmm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm5 = xmm3[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm5 = xmm5[0],zero,zero,zero,xmm5[1],zero,zero,zero,xmm5[2],zero,zero,zero,xmm5[3],zero,zero,zero,xmm5[4],zero,zero,zero,xmm5[5],zero,zero,zero,xmm5[6],zero,zero,zero,xmm5[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm5, %ymm10, %ymm5
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm3 = xmm3[0],zero,zero,zero,xmm3[1],zero,zero,zero,xmm3[2],zero,zero,zero,xmm3[3],zero,zero,zero,xmm3[4],zero,zero,zero,xmm3[5],zero,zero,zero,xmm3[6],zero,zero,zero,xmm3[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm3, %ymm9, %ymm3
+; AVXVNNI-NEXT: vpshufd {{.*#+}} xmm6 = xmm4[2,3,2,3]
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm6 = xmm6[0],zero,zero,zero,xmm6[1],zero,zero,zero,xmm6[2],zero,zero,zero,xmm6[3],zero,zero,zero,xmm6[4],zero,zero,zero,xmm6[5],zero,zero,zero,xmm6[6],zero,zero,zero,xmm6[7],zero,zero,zero
+; AVXVNNI-NEXT: vpmaddwd %ymm6, %ymm11, %ymm6
+; AVXVNNI-NEXT: vpmovzxbd {{.*#+}} ymm4 = xmm4[0],zero,zero,zero,xmm4[1],zero,zero,zero,xmm4[2],zero,zero,zero,xmm4[3],zero,zero,zero,xmm4[4],zero,zero,zero,xmm4[5],zero,zero,zero,xmm4[6],zero,zero,zero,xmm4[7],zero,zero,zero
+; AVXVNNI-NEXT: {vex} vpdpwssd %ymm4, %ymm2, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm6, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm5, %ymm0, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_umla_i8_v16i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm0, %ymm5
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm2, %ymm3
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm1, %ymm4
+; AVXVNNIINT8INT16-NEXT: vpdpbuud %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: vpdpbuud %ymm3, %ymm4, %ymm5
+; AVXVNNIINT8INT16-NEXT: vinsertf64x4 $1, %ymm5, %zmm0, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_umla_i8_v16i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpbuud %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.zext = zext <64 x i8> %a to <64 x i32>
+ %b.zext = zext <64 x i8> %b to <64 x i32>
+ %mul = mul nsw <64 x i32> %a.zext, %b.zext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> %acc, <64 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+; VNNI-INT16 i16 x i16 -> i32 SUMLA (sext x zext) tests: VPDPWSUD
+
+; Test 128-bit (xmm) SUMLA i16.
+define <4 x i32> @partial_reduce_sumla_i16_v4i32(<4 x i32> %acc, <8 x i16> %a, <8 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_sumla_i16_v4i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovsxwd %xmm1, %ymm1
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVX512VNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVX512VNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VNNI-NEXT: vzeroupper
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_sumla_i16_v4i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovsxwd %xmm1, %ymm1
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT: vzeroupper
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_sumla_i16_v4i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwsud %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_sumla_i16_v4i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwsud %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.sext = sext <8 x i16> %a to <8 x i32>
+ %b.zext = zext <8 x i16> %b to <8 x i32>
+ %mul = mul nsw <8 x i32> %a.sext, %b.zext
+ %res = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v8i32(<4 x i32> %acc, <8 x i32> %mul)
+ ret <4 x i32> %res
+}
+
+; Test 256-bit (ymm) SUMLA i16.
+define <8 x i32> @partial_reduce_sumla_i16_v8i32(<8 x i32> %acc, <16 x i16> %a, <16 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_sumla_i16_v8i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovsxwd %ymm1, %zmm1
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm2 = ymm2[0],zero,ymm2[1],zero,ymm2[2],zero,ymm2[3],zero,ymm2[4],zero,ymm2[5],zero,ymm2[6],zero,ymm2[7],zero,ymm2[8],zero,ymm2[9],zero,ymm2[10],zero,ymm2[11],zero,ymm2[12],zero,ymm2[13],zero,ymm2[14],zero,ymm2[15],zero
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vpaddd %ymm0, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_sumla_i16_v8i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovsxwd %xmm1, %ymm3
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpmovsxwd %xmm1, %ymm1
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm3, %ymm3
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_sumla_i16_v8i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwsud %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_sumla_i16_v8i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwsud %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.sext = sext <16 x i16> %a to <16 x i32>
+ %b.zext = zext <16 x i16> %b to <16 x i32>
+ %mul = mul nsw <16 x i32> %a.sext, %b.zext
+ %res = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v16i32(<8 x i32> %acc, <16 x i32> %mul)
+ ret <8 x i32> %res
+}
+
+; Test 512-bit (zmm) SUMLA i16.
+define <16 x i32> @partial_reduce_sumla_i16_v16i32(<16 x i32> %acc, <32 x i16> %a, <32 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_sumla_i16_v16i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovsxwd %ymm1, %zmm3
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm4 = ymm2[0],zero,ymm2[1],zero,ymm2[2],zero,ymm2[3],zero,ymm2[4],zero,ymm2[5],zero,ymm2[6],zero,ymm2[7],zero,ymm2[8],zero,ymm2[9],zero,ymm2[10],zero,ymm2[11],zero,ymm2[12],zero,ymm2[13],zero,ymm2[14],zero,ymm2[15],zero
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm2, %ymm2
+; AVX512VNNI-NEXT: vpmovsxwd %ymm1, %zmm1
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm2 = ymm2[0],zero,ymm2[1],zero,ymm2[2],zero,ymm2[3],zero,ymm2[4],zero,ymm2[5],zero,ymm2[6],zero,ymm2[7],zero,ymm2[8],zero,ymm2[9],zero,ymm2[10],zero,ymm2[11],zero,ymm2[12],zero,ymm2[13],zero,ymm2[14],zero,ymm2[15],zero
+; AVX512VNNI-NEXT: vpmulld %zmm4, %zmm3, %zmm3
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %zmm3, %zmm0, %zmm0
+; AVX512VNNI-NEXT: vpaddd %zmm1, %zmm0, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_sumla_i16_v16i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovsxwd %xmm2, %ymm6
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovsxwd %xmm2, %ymm2
+; AVXVNNI-NEXT: vpmovsxwd %xmm3, %ymm7
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm3
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm8 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm8, %ymm6, %ymm6
+; AVXVNNI-NEXT: vextracti128 $1, %ymm4, %xmm4
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm2, %ymm2
+; AVXVNNI-NEXT: vpmovsxwd %xmm3, %ymm3
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm5[0],zero,xmm5[1],zero,xmm5[2],zero,xmm5[3],zero,xmm5[4],zero,xmm5[5],zero,xmm5[6],zero,xmm5[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm7, %ymm4
+; AVXVNNI-NEXT: vextracti128 $1, %ymm5, %xmm5
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm5 = xmm5[0],zero,xmm5[1],zero,xmm5[2],zero,xmm5[3],zero,xmm5[4],zero,xmm5[5],zero,xmm5[6],zero,xmm5[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm5, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpaddd %ymm6, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm2, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm4, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm1, %ymm1
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_sumla_i16_v16i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm0, %ymm5
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm2, %ymm3
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm1, %ymm4
+; AVXVNNIINT8INT16-NEXT: vpdpwsud %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: vpdpwsud %ymm3, %ymm4, %ymm5
+; AVXVNNIINT8INT16-NEXT: vinsertf64x4 $1, %ymm5, %zmm0, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_sumla_i16_v16i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwsud %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.sext = sext <32 x i16> %a to <32 x i32>
+ %b.zext = zext <32 x i16> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.sext, %b.zext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> %acc, <32 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+; VNNI-INT16 i16 x i16 -> i32 UMLA (zext x zext) tests: VPDPWUUD
+
+; Test 128-bit (xmm) UMLA i16.
+define <4 x i32> @partial_reduce_umla_i16_v4i32(<4 x i32> %acc, <8 x i16> %a, <8 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_umla_i16_v4i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} ymm1 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVX512VNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVX512VNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVX512VNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVX512VNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVX512VNNI-NEXT: vzeroupper
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_umla_i16_v4i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm1 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %xmm1, %xmm0, %xmm0
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpaddd %xmm0, %xmm1, %xmm0
+; AVXVNNI-NEXT: vzeroupper
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_umla_i16_v4i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwuud %xmm2, %xmm1, %xmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_umla_i16_v4i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwuud %xmm2, %xmm1, %xmm0
+; AVX10-NEXT: retq
+ %a.zext = zext <8 x i16> %a to <8 x i32>
+ %b.zext = zext <8 x i16> %b to <8 x i32>
+ %mul = mul nsw <8 x i32> %a.zext, %b.zext
+ %res = call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v8i32(<4 x i32> %acc, <8 x i32> %mul)
+ ret <4 x i32> %res
+}
+
+; Test 256-bit (ymm) UMLA i16.
+define <8 x i32> @partial_reduce_umla_i16_v8i32(<8 x i32> %acc, <16 x i16> %a, <16 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_umla_i16_v8i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm1 = ymm1[0],zero,ymm1[1],zero,ymm1[2],zero,ymm1[3],zero,ymm1[4],zero,ymm1[5],zero,ymm1[6],zero,ymm1[7],zero,ymm1[8],zero,ymm1[9],zero,ymm1[10],zero,ymm1[11],zero,ymm1[12],zero,ymm1[13],zero,ymm1[14],zero,ymm1[15],zero
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm2 = ymm2[0],zero,ymm2[1],zero,ymm2[2],zero,ymm2[3],zero,ymm2[4],zero,ymm2[5],zero,ymm2[6],zero,ymm2[7],zero,ymm2[8],zero,ymm2[9],zero,ymm2[10],zero,ymm2[11],zero,ymm2[12],zero,ymm2[13],zero,ymm2[14],zero,ymm2[15],zero
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vpaddd %ymm0, %ymm1, %ymm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_umla_i16_v8i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm3 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVXVNNI-NEXT: vextracti128 $1, %ymm1, %xmm1
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm1 = xmm1[0],zero,xmm1[1],zero,xmm1[2],zero,xmm1[3],zero,xmm1[4],zero,xmm1[5],zero,xmm1[6],zero,xmm1[7],zero
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm3, %ymm3
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm2, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm1, %ymm0, %ymm0
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_umla_i16_v8i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vpdpwuud %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_umla_i16_v8i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwuud %ymm2, %ymm1, %ymm0
+; AVX10-NEXT: retq
+ %a.zext = zext <16 x i16> %a to <16 x i32>
+ %b.zext = zext <16 x i16> %b to <16 x i32>
+ %mul = mul nsw <16 x i32> %a.zext, %b.zext
+ %res = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v16i32(<8 x i32> %acc, <16 x i32> %mul)
+ ret <8 x i32> %res
+}
+
+; Test 512-bit (zmm) UMLA i16.
+define <16 x i32> @partial_reduce_umla_i16_v16i32(<16 x i32> %acc, <32 x i16> %a, <32 x i16> %b) {
+; AVX512VNNI-LABEL: partial_reduce_umla_i16_v16i32:
+; AVX512VNNI: # %bb.0:
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm3 = ymm1[0],zero,ymm1[1],zero,ymm1[2],zero,ymm1[3],zero,ymm1[4],zero,ymm1[5],zero,ymm1[6],zero,ymm1[7],zero,ymm1[8],zero,ymm1[9],zero,ymm1[10],zero,ymm1[11],zero,ymm1[12],zero,ymm1[13],zero,ymm1[14],zero,ymm1[15],zero
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm1, %ymm1
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm4 = ymm2[0],zero,ymm2[1],zero,ymm2[2],zero,ymm2[3],zero,ymm2[4],zero,ymm2[5],zero,ymm2[6],zero,ymm2[7],zero,ymm2[8],zero,ymm2[9],zero,ymm2[10],zero,ymm2[11],zero,ymm2[12],zero,ymm2[13],zero,ymm2[14],zero,ymm2[15],zero
+; AVX512VNNI-NEXT: vextracti64x4 $1, %zmm2, %ymm2
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm1 = ymm1[0],zero,ymm1[1],zero,ymm1[2],zero,ymm1[3],zero,ymm1[4],zero,ymm1[5],zero,ymm1[6],zero,ymm1[7],zero,ymm1[8],zero,ymm1[9],zero,ymm1[10],zero,ymm1[11],zero,ymm1[12],zero,ymm1[13],zero,ymm1[14],zero,ymm1[15],zero
+; AVX512VNNI-NEXT: vpmovzxwd {{.*#+}} zmm2 = ymm2[0],zero,ymm2[1],zero,ymm2[2],zero,ymm2[3],zero,ymm2[4],zero,ymm2[5],zero,ymm2[6],zero,ymm2[7],zero,ymm2[8],zero,ymm2[9],zero,ymm2[10],zero,ymm2[11],zero,ymm2[12],zero,ymm2[13],zero,ymm2[14],zero,ymm2[15],zero
+; AVX512VNNI-NEXT: vpmulld %zmm4, %zmm3, %zmm3
+; AVX512VNNI-NEXT: vpmulld %zmm2, %zmm1, %zmm1
+; AVX512VNNI-NEXT: vpaddd %zmm3, %zmm0, %zmm0
+; AVX512VNNI-NEXT: vpaddd %zmm1, %zmm0, %zmm0
+; AVX512VNNI-NEXT: retq
+;
+; AVXVNNI-LABEL: partial_reduce_umla_i16_v16i32:
+; AVXVNNI: # %bb.0:
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm6 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vextracti128 $1, %ymm2, %xmm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm2 = xmm2[0],zero,xmm2[1],zero,xmm2[2],zero,xmm2[3],zero,xmm2[4],zero,xmm2[5],zero,xmm2[6],zero,xmm2[7],zero
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm7 = xmm3[0],zero,xmm3[1],zero,xmm3[2],zero,xmm3[3],zero,xmm3[4],zero,xmm3[5],zero,xmm3[6],zero,xmm3[7],zero
+; AVXVNNI-NEXT: vextracti128 $1, %ymm3, %xmm3
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm8 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm8, %ymm6, %ymm6
+; AVXVNNI-NEXT: vextracti128 $1, %ymm4, %xmm4
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm4[0],zero,xmm4[1],zero,xmm4[2],zero,xmm4[3],zero,xmm4[4],zero,xmm4[5],zero,xmm4[6],zero,xmm4[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm2, %ymm2
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm3 = xmm3[0],zero,xmm3[1],zero,xmm3[2],zero,xmm3[3],zero,xmm3[4],zero,xmm3[5],zero,xmm3[6],zero,xmm3[7],zero
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm4 = xmm5[0],zero,xmm5[1],zero,xmm5[2],zero,xmm5[3],zero,xmm5[4],zero,xmm5[5],zero,xmm5[6],zero,xmm5[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm4, %ymm7, %ymm4
+; AVXVNNI-NEXT: vextracti128 $1, %ymm5, %xmm5
+; AVXVNNI-NEXT: vpmovzxwd {{.*#+}} ymm5 = xmm5[0],zero,xmm5[1],zero,xmm5[2],zero,xmm5[3],zero,xmm5[4],zero,xmm5[5],zero,xmm5[6],zero,xmm5[7],zero
+; AVXVNNI-NEXT: vpmulld %ymm5, %ymm3, %ymm3
+; AVXVNNI-NEXT: vpaddd %ymm6, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm2, %ymm0, %ymm0
+; AVXVNNI-NEXT: vpaddd %ymm4, %ymm1, %ymm1
+; AVXVNNI-NEXT: vpaddd %ymm3, %ymm1, %ymm1
+; AVXVNNI-NEXT: retq
+;
+; AVXVNNIINT8INT16-LABEL: partial_reduce_umla_i16_v16i32:
+; AVXVNNIINT8INT16: # %bb.0:
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm0, %ymm5
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm2, %ymm3
+; AVXVNNIINT8INT16-NEXT: vextractf64x4 $1, %zmm1, %ymm4
+; AVXVNNIINT8INT16-NEXT: vpdpwuud %ymm2, %ymm1, %ymm0
+; AVXVNNIINT8INT16-NEXT: vpdpwuud %ymm3, %ymm4, %ymm5
+; AVXVNNIINT8INT16-NEXT: vinsertf64x4 $1, %ymm5, %zmm0, %zmm0
+; AVXVNNIINT8INT16-NEXT: retq
+;
+; AVX10-LABEL: partial_reduce_umla_i16_v16i32:
+; AVX10: # %bb.0:
+; AVX10-NEXT: vpdpwuud %zmm2, %zmm1, %zmm0
+; AVX10-NEXT: retq
+ %a.zext = zext <32 x i16> %a to <32 x i32>
+ %b.zext = zext <32 x i16> %b to <32 x i32>
+ %mul = mul nsw <32 x i32> %a.zext, %b.zext
+ %res = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> %acc, <32 x i32> %mul)
+ ret <16 x i32> %res
+}
+
+declare <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v8i32(<4 x i32>, <8 x i32>)
+declare <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v16i32(<8 x i32>, <16 x i32>)
+declare <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32>, <32 x i32>)
+declare <4 x float> @llvm.vector.partial.reduce.fadd.v4f32.v8f32(<4 x float>, <8 x float>)
+declare <8 x float> @llvm.vector.partial.reduce.fadd.v8f32.v16f32(<8 x float>, <16 x float>)
+declare <16 x float> @llvm.vector.partial.reduce.fadd.v16f32.v32f32(<16 x float>, <32 x float>)
diff --git a/llvm/test/CodeGen/X86/partial-reduce-small-vf.ll b/llvm/test/CodeGen/X86/partial-reduce-small-vf.ll
new file mode 100644
index 0000000000000..2468dd5aa9a46
--- /dev/null
+++ b/llvm/test/CodeGen/X86/partial-reduce-small-vf.ll
@@ -0,0 +1,38 @@
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=znver4 -verify-machineinstrs < %s | FileCheck %s
+
+; Unsupported small partial-reduction shapes should expand or widen safely
+; instead of crashing during type legalization.
+
+define <2 x i32> @partial_reduce_sumla_i8_v2i32(<2 x i32> %acc, <8 x i8> %a, <8 x i8> %b) {
+; CHECK-LABEL: partial_reduce_sumla_i8_v2i32:
+; CHECK: retq
+ %a.zext = zext <8 x i8> %a to <8 x i32>
+ %b.sext = sext <8 x i8> %b to <8 x i32>
+ %mul = mul nsw <8 x i32> %a.zext, %b.sext
+ %res = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v8i32(<2 x i32> %acc, <8 x i32> %mul)
+ ret <2 x i32> %res
+}
+
+define <2 x i32> @partial_reduce_smla_i16_v2i32(<2 x i32> %acc, <4 x i16> %a, <4 x i16> %b) {
+; CHECK-LABEL: partial_reduce_smla_i16_v2i32:
+; CHECK: retq
+ %a.sext = sext <4 x i16> %a to <4 x i32>
+ %b.sext = sext <4 x i16> %b to <4 x i32>
+ %mul = mul nsw <4 x i32> %a.sext, %b.sext
+ %res = call <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(<2 x i32> %acc, <4 x i32> %mul)
+ ret <2 x i32> %res
+}
+
+define <2 x float> @partial_reduce_fmla_bf16_v2f32(<2 x float> %acc, <4 x bfloat> %a, <4 x bfloat> %b) {
+; CHECK-LABEL: partial_reduce_fmla_bf16_v2f32:
+; CHECK: retq
+ %a.ext = fpext <4 x bfloat> %a to <4 x float>
+ %b.ext = fpext <4 x bfloat> %b to <4 x float>
+ %mul = fmul <4 x float> %a.ext, %b.ext
+ %res = call <2 x float> @llvm.vector.partial.reduce.fadd.v2f32.v4f32(<2 x float> %acc, <4 x float> %mul)
+ ret <2 x float> %res
+}
+
+declare <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v8i32(<2 x i32>, <8 x i32>)
+declare <2 x i32> @llvm.vector.partial.reduce.add.v2i32.v4i32(<2 x i32>, <4 x i32>)
+declare <2 x float> @llvm.vector.partial.reduce.fadd.v2f32.v4f32(<2 x float>, <4 x float>)
diff --git a/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-cost-model.ll b/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-cost-model.ll
new file mode 100644
index 0000000000000..fef6aaf6dace7
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-cost-model.ll
@@ -0,0 +1,107 @@
+; RUN: opt < %s -passes=loop-vectorize -force-vector-interleave=1 -enable-epilogue-vectorization=false -debug-only=loop-vectorize --disable-output 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+; A partial reduction cannot consume an extended value that still feeds another
+; reduction. The cost model must keep the recurrence width at i32.
+define i32 @dot_extra_extend_user_two_reductions(ptr readonly %a, ptr readonly %b) #0 {
+; CHECK-LABEL: LV: Checking a loop in 'dot_extra_extend_user_two_reductions'
+; CHECK-NOT: LV: Partial reductions narrow WidestType
+; CHECK: LV: The Smallest and Widest types: 8 / 32 bits.
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %aux = phi i32 [ 0, %entry ], [ %aux.next, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = sext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %aux.next = add nsw i32 %aux, %ext.a
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ %r = add nsw i32 %add, %aux.next
+ ret i32 %r
+}
+
+; Mixed-scale reduction chains are not representable as partial reductions.
+; The cost model must not narrow to the first i8 link and ignore the later i16
+; link.
+define i32 @dot_mixed_scale_no_partial_reduce(ptr readonly %a, ptr readonly %b,
+ ptr readonly %c, ptr readonly %d) #1 {
+; CHECK-LABEL: LV: Checking a loop in 'dot_mixed_scale_no_partial_reduce'
+; CHECK-NOT: LV: Partial reductions narrow WidestType
+; CHECK: LV: The Smallest and Widest types: 8 / 32 bits.
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add2, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = sext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul8 = mul nsw i32 %ext.a, %ext.b
+ %add1 = add nsw i32 %sum, %mul8
+ %gep.c = getelementptr inbounds i16, ptr %c, i64 %iv
+ %load.c = load i16, ptr %gep.c, align 2
+ %ext.c = sext i16 %load.c to i32
+ %gep.d = getelementptr inbounds i16, ptr %d, i64 %iv
+ %load.d = load i16, ptr %gep.d, align 2
+ %ext.d = sext i16 %load.d to i32
+ %mul16 = mul nsw i32 %ext.c, %ext.d
+ %add2 = add nsw i32 %add1, %mul16
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add2
+}
+
+; Partial reductions can narrow the effective recurrence width only when the
+; widened recurrence type is not also needed by unrelated loop operations.
+define i32 @dot_with_unrelated_i32_store(ptr readonly %a, ptr readonly %b,
+ ptr %c) #1 {
+; CHECK-LABEL: LV: Checking a loop in 'dot_with_unrelated_i32_store'
+; CHECK-NOT: LV: Partial reductions narrow WidestType
+; CHECK: LV: The Smallest and Widest types: 8 / 32 bits.
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %gep.c = getelementptr inbounds i32, ptr %c, i64 %iv
+ %load.c = load i32, ptr %gep.c, align 4
+ %inc.c = add nsw i32 %load.c, 1
+ store i32 %inc.c, ptr %gep.c, align 4
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+attributes #0 = { "target-cpu"="alderlake" "target-features"="+avx2,+avxvnniint8,-avx512f,-avx512vnni" }
+attributes #1 = { "target-cpu"="alderlake" "target-features"="+avx2,+avxvnni,+avxvnniint8,-avx512f,-avx512vnni" }
diff --git a/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-dot-product.ll b/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-dot-product.ll
new file mode 100644
index 0000000000000..ea3f1b0f6b068
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-dot-product.ll
@@ -0,0 +1,653 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals none --filter-out-after "^scalar.ph:" --version 4
+; RUN: opt -passes=loop-vectorize -force-vector-interleave=1 -enable-epilogue-vectorization=false -S < %s | FileCheck %s --check-prefixes=CHECK
+; RUN: opt -passes=loop-vectorize -S < %s | FileCheck %s --check-prefixes=RUNTIME
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Test u8×i8 -> i32 dot product is vectorized using partial reductions
+; on AVX512-VNNI targets, producing llvm.vector.partial.reduce.add with
+; <64 x i8> inputs and <16 x i32> accumulator (512-bit zmm).
+define i32 @dot_product_i8(ptr readonly %a, ptr readonly %b) #0 {
+; CHECK-LABEL: define i32 @dot_product_i8(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <64 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <64 x i8>, ptr [[TMP1]], align 1
+; CHECK-NEXT: [[TMP2:%.*]] = zext <64 x i8> [[WIDE_LOAD]] to <64 x i32>
+; CHECK-NEXT: [[TMP3:%.*]] = sext <64 x i8> [[WIDE_LOAD1]] to <64 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <64 x i32> [[TMP2]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> [[VEC_PHI]], <64 x i32> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 64
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP6:%.*]] = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP6]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test u8×i8 -> i32 dot product with a runtime trip count. This covers the
+; issue #187779 shape where both the main vector loop and vector epilogue should
+; keep using partial reductions.
+define i32 @dot_product_i8_runtime_n(ptr readonly %a, ptr readonly %b, i64 %n) #0 {
+; RUNTIME-LABEL: define i32 @dot_product_i8_runtime_n(
+; RUNTIME: vector.body:
+; RUNTIME: call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32
+; RUNTIME: middle.block:
+; RUNTIME: vec.epilog.vector.body:
+; RUNTIME: call <4 x i32> @llvm.vector.partial.reduce.add.v4i32.v16i32
+; RUNTIME: vec.epilog.middle.block:
+;
+entry:
+ %cmp = icmp eq i64 %n, 0
+ br i1 %cmp, label %exit, label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, %n
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ %res = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ ret i32 %res
+}
+
+; Test that without VNNI, no partial reduction is generated.
+define i32 @dot_product_i8_no_vnni(ptr readonly %a, ptr readonly %b) #1 {
+; CHECK-LABEL: define i32 @dot_product_i8_no_vnni(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <8 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[TMP5:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <8 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = zext <8 x i8> [[WIDE_LOAD]] to <8 x i32>
+; CHECK-NEXT: [[TMP2:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <8 x i8>, ptr [[TMP2]], align 1
+; CHECK-NEXT: [[TMP3:%.*]] = sext <8 x i8> [[WIDE_LOAD1]] to <8 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <8 x i32> [[TMP1]], [[TMP3]]
+; CHECK-NEXT: [[TMP5]] = add <8 x i32> [[VEC_PHI]], [[TMP4]]
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 8
+; CHECK-NEXT: [[TMP6:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP6]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP7:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[TMP5]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP7]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test i16×i16 -> i32 dot product is vectorized using partial reductions
+; on AVX512-VNNI targets, producing llvm.vector.partial.reduce.add with
+; <32 x i16> inputs and <16 x i32> accumulator (512-bit zmm).
+define i32 @dot_product_i16(ptr readonly %a, ptr readonly %b) #0 {
+; CHECK-LABEL: define i32 @dot_product_i16(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i16, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <32 x i16>, ptr [[TMP0]], align 2
+; CHECK-NEXT: [[TMP2:%.*]] = getelementptr inbounds i16, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <32 x i16>, ptr [[TMP2]], align 2
+; CHECK-NEXT: [[TMP5:%.*]] = sext <32 x i16> [[WIDE_LOAD]] to <32 x i32>
+; CHECK-NEXT: [[TMP3:%.*]] = sext <32 x i16> [[WIDE_LOAD1]] to <32 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <32 x i32> [[TMP5]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v32i32(<16 x i32> [[VEC_PHI]], <32 x i32> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 32
+; CHECK-NEXT: [[TMP6:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP6]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP7:%.*]] = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP7]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i16, ptr %a, i64 %iv
+ %load.a = load i16, ptr %gep.a, align 2
+ %ext.a = sext i16 %load.a to i32
+ %gep.b = getelementptr inbounds i16, ptr %b, i64 %iv
+ %load.b = load i16, ptr %gep.b, align 2
+ %ext.b = sext i16 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test bf16×bf16 -> f32 dot product is vectorized using partial reductions
+; on AVX512-BF16 targets, producing llvm.vector.partial.reduce.fadd with
+; <32 x bfloat> inputs and <16 x float> accumulator (512-bit zmm).
+define float @dot_product_bf16(ptr readonly %a, ptr readonly %b) #2 {
+; CHECK-LABEL: define float @dot_product_bf16(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x float> [ <float 0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds bfloat, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <32 x bfloat>, ptr [[TMP0]], align 2
+; CHECK-NEXT: [[TMP2:%.*]] = getelementptr inbounds bfloat, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <32 x bfloat>, ptr [[TMP2]], align 2
+; CHECK-NEXT: [[TMP5:%.*]] = fpext <32 x bfloat> [[WIDE_LOAD]] to <32 x float>
+; CHECK-NEXT: [[TMP3:%.*]] = fpext <32 x bfloat> [[WIDE_LOAD1]] to <32 x float>
+; CHECK-NEXT: [[TMP4:%.*]] = fmul contract <32 x float> [[TMP5]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call reassoc contract <16 x float> @llvm.vector.partial.reduce.fadd.v16f32.v32f32(<16 x float> [[VEC_PHI]], <32 x float> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 32
+; CHECK-NEXT: [[TMP6:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP6]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP5:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP7:%.*]] = call reassoc contract float @llvm.vector.reduce.fadd.v16f32(float -0.000000e+00, <16 x float> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret float [[TMP7]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi float [ 0.0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds bfloat, ptr %a, i64 %iv
+ %load.a = load bfloat, ptr %gep.a, align 2
+ %ext.a = fpext bfloat %load.a to float
+ %gep.b = getelementptr inbounds bfloat, ptr %b, i64 %iv
+ %load.b = load bfloat, ptr %gep.b, align 2
+ %ext.b = fpext bfloat %load.b to float
+ %mul = fmul contract float %ext.a, %ext.b
+ %add = fadd reassoc contract float %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret float %add
+}
+
+; Test sext*sext i8 -> i32 dot product with VNNI-INT8: should use partial reduce.
+define i32 @dot_product_s8s8_vnniint8(ptr readonly %a, ptr readonly %b) #3 {
+; CHECK-LABEL: define i32 @dot_product_s8s8_vnniint8(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR3:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <64 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <64 x i8>, ptr [[TMP1]], align 1
+; CHECK-NEXT: [[TMP2:%.*]] = sext <64 x i8> [[WIDE_LOAD]] to <64 x i32>
+; CHECK-NEXT: [[TMP3:%.*]] = sext <64 x i8> [[WIDE_LOAD1]] to <64 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <64 x i32> [[TMP2]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> [[VEC_PHI]], <64 x i32> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 64
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP6:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP6:%.*]] = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP6]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = sext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test zext*zext i8 -> i32 dot product with VNNI-INT8: should use partial reduce.
+define i32 @dot_product_u8u8_vnniint8(ptr readonly %a, ptr readonly %b) #3 {
+; CHECK-LABEL: define i32 @dot_product_u8u8_vnniint8(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR3]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <64 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <64 x i8>, ptr [[TMP1]], align 1
+; CHECK-NEXT: [[TMP2:%.*]] = zext <64 x i8> [[WIDE_LOAD]] to <64 x i32>
+; CHECK-NEXT: [[TMP3:%.*]] = zext <64 x i8> [[WIDE_LOAD1]] to <64 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <64 x i32> [[TMP2]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call <16 x i32> @llvm.vector.partial.reduce.add.v16i32.v64i32(<16 x i32> [[VEC_PHI]], <64 x i32> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 64
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP7:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP6:%.*]] = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP6]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = zext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test sext*sext i8 -> i32 on plain VNNI (no VNNI-INT8): should NOT use partial reduce.
+define i32 @dot_product_s8s8_no_vnniint8(ptr readonly %a, ptr readonly %b) #0 {
+; CHECK-LABEL: define i32 @dot_product_s8s8_no_vnniint8(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[TMP5:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <16 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = sext <16 x i8> [[WIDE_LOAD]] to <16 x i32>
+; CHECK-NEXT: [[TMP2:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <16 x i8>, ptr [[TMP2]], align 1
+; CHECK-NEXT: [[TMP3:%.*]] = sext <16 x i8> [[WIDE_LOAD1]] to <16 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <16 x i32> [[TMP1]], [[TMP3]]
+; CHECK-NEXT: [[TMP5]] = add <16 x i32> [[VEC_PHI]], [[TMP4]]
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 16
+; CHECK-NEXT: [[TMP6:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP6]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP8:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP7:%.*]] = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> [[TMP5]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP7]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = sext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test zext*zext i8 -> i32 on plain VNNI (no VNNI-INT8): should NOT use partial reduce.
+define i32 @dot_product_u8u8_no_vnniint8(ptr readonly %a, ptr readonly %b) #0 {
+; CHECK-LABEL: define i32 @dot_product_u8u8_no_vnniint8(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <16 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[TMP5:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <16 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = zext <16 x i8> [[WIDE_LOAD]] to <16 x i32>
+; CHECK-NEXT: [[TMP2:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <16 x i8>, ptr [[TMP2]], align 1
+; CHECK-NEXT: [[TMP3:%.*]] = zext <16 x i8> [[WIDE_LOAD1]] to <16 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <16 x i32> [[TMP1]], [[TMP3]]
+; CHECK-NEXT: [[TMP5]] = add <16 x i32> [[VEC_PHI]], [[TMP4]]
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 16
+; CHECK-NEXT: [[TMP6:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP6]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP9:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP7:%.*]] = call i32 @llvm.vector.reduce.add.v16i32(<16 x i32> [[TMP5]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP7]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = zext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test u8×i8 -> i32 dot product on AVXVNNI (non-AVX512): should use partial
+; reduction with 256-bit ymm vectors (VF=32, <32 x i8> -> <8 x i32>).
+define i32 @dot_product_i8_avxvnni(ptr readonly %a, ptr readonly %b) #4 {
+; CHECK-LABEL: define i32 @dot_product_i8_avxvnni(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR4:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <8 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <32 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <32 x i8>, ptr [[TMP1]], align 1
+; CHECK-NEXT: [[TMP2:%.*]] = zext <32 x i8> [[WIDE_LOAD]] to <32 x i32>
+; CHECK-NEXT: [[TMP3:%.*]] = sext <32 x i8> [[WIDE_LOAD1]] to <32 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <32 x i32> [[TMP2]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v32i32(<8 x i32> [[VEC_PHI]], <32 x i32> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 32
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP10:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP6:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP6]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Test bf16 dot product WITHOUT avx512bf16: should NOT use partial reduce.
+define float @dot_product_bf16_no_bf16(ptr readonly %a, ptr readonly %b) #5 {
+; CHECK-LABEL: define float @dot_product_bf16_no_bf16(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR5:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <8 x float> [ <float 0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>, [[VECTOR_PH]] ], [ [[TMP7:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds bfloat, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <8 x bfloat>, ptr [[TMP0]], align 2
+; CHECK-NEXT: [[TMP2:%.*]] = fpext <8 x bfloat> [[WIDE_LOAD]] to <8 x float>
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds bfloat, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <8 x bfloat>, ptr [[TMP1]], align 2
+; CHECK-NEXT: [[TMP3:%.*]] = fpext <8 x bfloat> [[WIDE_LOAD1]] to <8 x float>
+; CHECK-NEXT: [[TMP4:%.*]] = fmul contract <8 x float> [[TMP2]], [[TMP3]]
+; CHECK-NEXT: [[TMP7]] = fadd reassoc contract <8 x float> [[VEC_PHI]], [[TMP4]]
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 8
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP11:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP6:%.*]] = call reassoc contract float @llvm.vector.reduce.fadd.v8f32(float -0.000000e+00, <8 x float> [[TMP7]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret float [[TMP6]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi float [ 0.0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds bfloat, ptr %a, i64 %iv
+ %load.a = load bfloat, ptr %gep.a, align 2
+ %ext.a = fpext bfloat %load.a to float
+ %gep.b = getelementptr inbounds bfloat, ptr %b, i64 %iv
+ %load.b = load bfloat, ptr %gep.b, align 2
+ %ext.b = fpext bfloat %load.b to float
+ %mul = fmul contract float %ext.a, %ext.b
+ %add = fadd reassoc contract float %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret float %add
+}
+
+; Test bf16 -> f32 summation stays as a regular reduction even when bf16
+; partial reductions are available for dot products.
+define float @sum_bf16_no_partial_reduce(ptr readonly %a) #2 {
+; CHECK-LABEL: define float @sum_bf16_no_partial_reduce(
+; CHECK-SAME: ptr readonly [[A:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <4 x float> [ <float 0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>, [[VECTOR_PH]] ], [ [[VEC_ADD:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds bfloat, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <4 x bfloat>, ptr [[TMP0]], align 2
+; CHECK-NEXT: [[TMP1:%.*]] = fpext <4 x bfloat> [[WIDE_LOAD]] to <4 x float>
+; CHECK-NEXT: [[VEC_ADD]] = fadd reassoc contract <4 x float> [[VEC_PHI]], [[TMP1]]
+; CHECK-NOT: @llvm.vector.partial.reduce.fadd
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
+; CHECK-NEXT: [[TMP2:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP2]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP12:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP3:%.*]] = call reassoc contract float @llvm.vector.reduce.fadd.v4f32(float -0.000000e+00, <4 x float> [[VEC_ADD]])
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret float [[TMP3]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi float [ 0.0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds bfloat, ptr %a, i64 %iv
+ %load.a = load bfloat, ptr %gep.a, align 2
+ %ext.a = fpext bfloat %load.a to float
+ %add = fadd reassoc contract float %sum, %ext.a
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret float %add
+}
+
+; Test subtract dot-product on a 256-bit VNNI-INT8 target still forms a
+; partial reduction and applies the final subtract in the middle block.
+define i32 @dot_sub_s8s8_vnniint8_avx2(ptr readonly %a, ptr readonly %b) #6 {
+; CHECK-LABEL: define i32 @dot_sub_s8s8_vnniint8_avx2(
+; CHECK-SAME: ptr readonly [[A:%.*]], ptr readonly [[B:%.*]]) #[[ATTR6:[0-9]+]] {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: br label [[VECTOR_PH:%.*]]
+; CHECK: vector.ph:
+; CHECK-NEXT: br label [[VECTOR_BODY:%.*]]
+; CHECK: vector.body:
+; CHECK-NEXT: [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[VEC_PHI:%.*]] = phi <8 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[PARTIAL_REDUCE:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT: [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[A]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD:%.*]] = load <32 x i8>, ptr [[TMP0]], align 1
+; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[B]], i64 [[INDEX]]
+; CHECK-NEXT: [[WIDE_LOAD1:%.*]] = load <32 x i8>, ptr [[TMP1]], align 1
+; CHECK-NEXT: [[TMP2:%.*]] = sext <32 x i8> [[WIDE_LOAD]] to <32 x i32>
+; CHECK-NEXT: [[TMP3:%.*]] = sext <32 x i8> [[WIDE_LOAD1]] to <32 x i32>
+; CHECK-NEXT: [[TMP4:%.*]] = mul nsw <32 x i32> [[TMP2]], [[TMP3]]
+; CHECK-NEXT: [[PARTIAL_REDUCE]] = call <8 x i32> @llvm.vector.partial.reduce.add.v8i32.v32i32(<8 x i32> [[VEC_PHI]], <32 x i32> [[TMP4]])
+; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 32
+; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1024
+; CHECK-NEXT: br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP13:![0-9]+]]
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP6:%.*]] = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> [[PARTIAL_REDUCE]])
+; CHECK-NEXT: [[TMP7:%.*]] = sub i32 0, [[TMP6]]
+; CHECK-NEXT: br label [[EXIT:%.*]]
+; CHECK: exit:
+; CHECK-NEXT: ret i32 [[TMP7]]
+;
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %sub, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = sext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %sub = sub nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %sub
+}
+
+attributes #0 = { "target-cpu"="znver4" "target-features"="+avx512vnni,+avx512bw,+avx512f,+avx512vl" }
+attributes #1 = { "target-cpu"="haswell" "target-features"="+avx2,-avx512vnni,-avxvnni" }
+attributes #2 = { "target-cpu"="znver4" "target-features"="+avx512vnni,+avx512bw,+avx512f,+avx512vl,+avx512bf16" }
+attributes #3 = { "target-cpu"="znver4" "target-features"="+avx512vnni,+avx512bw,+avx512f,+avx512vl,+avxvnniint8" }
+attributes #4 = { "target-cpu"="alderlake" "target-features"="+avxvnni,+avx2,-avx512f,-avx512vnni" }
+attributes #5 = { "target-cpu"="skylake-avx512" "target-features"="+avx512f,+avx512bw,+avx512vl,-avx512bf16" }
+attributes #6 = { "target-cpu"="alderlake" "target-features"="+avx2,+avxvnniint8,-avx512f,-avx512vnni" }
diff --git a/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-small-vf.ll b/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-small-vf.ll
new file mode 100644
index 0000000000000..ffcf8758fb7cd
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/X86/partial-reduce-small-vf.ll
@@ -0,0 +1,96 @@
+; RUN: opt -passes=loop-vectorize -force-vector-width=8 -force-vector-interleave=1 -enable-epilogue-vectorization=false -S < %s | FileCheck %s --check-prefix=I8
+; RUN: opt -passes=loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -enable-epilogue-vectorization=false -S < %s | FileCheck %s --check-prefix=I16
+; RUN: opt -passes=loop-vectorize -force-vector-width=4 -force-vector-interleave=1 -enable-epilogue-vectorization=false -S < %s | FileCheck %s --check-prefix=BF16
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Do not form a partial reduction when the forced VF would create a v2i32
+; accumulator. X86 only registers native VNNI partial-reduction actions with at
+; least four accumulator lanes.
+define i32 @dot_product_i8_forced_vf8(ptr readonly %a, ptr readonly %b) #0 {
+; I8-LABEL: define i32 @dot_product_i8_forced_vf8(
+; I8: vector.body:
+; I8-NOT: llvm.vector.partial.reduce
+; I8: middle.block:
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i8, ptr %a, i64 %iv
+ %load.a = load i8, ptr %gep.a, align 1
+ %ext.a = zext i8 %load.a to i32
+ %gep.b = getelementptr inbounds i8, ptr %b, i64 %iv
+ %load.b = load i8, ptr %gep.b, align 1
+ %ext.b = sext i8 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Do not form a partial reduction when the forced VF would create a v2i32
+; accumulator for i16 dot products either.
+define i32 @dot_product_i16_forced_vf4(ptr readonly %a, ptr readonly %b) #0 {
+; I16-LABEL: define i32 @dot_product_i16_forced_vf4(
+; I16: vector.body:
+; I16-NOT: llvm.vector.partial.reduce
+; I16: middle.block:
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds i16, ptr %a, i64 %iv
+ %load.a = load i16, ptr %gep.a, align 2
+ %ext.a = sext i16 %load.a to i32
+ %gep.b = getelementptr inbounds i16, ptr %b, i64 %iv
+ %load.b = load i16, ptr %gep.b, align 2
+ %ext.b = sext i16 %load.b to i32
+ %mul = mul nsw i32 %ext.a, %ext.b
+ %add = add nsw i32 %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret i32 %add
+}
+
+; Likewise for bf16, avoid a v2f32 partial reduction.
+define float @dot_product_bf16_forced_vf4(ptr readonly %a, ptr readonly %b) #1 {
+; BF16-LABEL: define float @dot_product_bf16_forced_vf4(
+; BF16: vector.body:
+; BF16-NOT: llvm.vector.partial.reduce
+; BF16: middle.block:
+entry:
+ br label %for.body
+
+for.body:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+ %sum = phi float [ 0.0, %entry ], [ %add, %for.body ]
+ %gep.a = getelementptr inbounds bfloat, ptr %a, i64 %iv
+ %load.a = load bfloat, ptr %gep.a, align 2
+ %ext.a = fpext bfloat %load.a to float
+ %gep.b = getelementptr inbounds bfloat, ptr %b, i64 %iv
+ %load.b = load bfloat, ptr %gep.b, align 2
+ %ext.b = fpext bfloat %load.b to float
+ %mul = fmul contract float %ext.a, %ext.b
+ %add = fadd reassoc contract float %sum, %mul
+ %iv.next = add nuw nsw i64 %iv, 1
+ %exitcond = icmp eq i64 %iv.next, 1024
+ br i1 %exitcond, label %exit, label %for.body
+
+exit:
+ ret float %add
+}
+
+attributes #0 = { "target-cpu"="znver4" }
+attributes #1 = { "target-cpu"="znver4" "target-features"="+avx512bf16,+avx512vl" }
More information about the llvm-commits
mailing list