[llvm] dfa7876 - [RISCV] Cost i64 accumulator for Zvdot4a8i partial reductions
via llvm-commits
llvm-commits at lists.llvm.org
Mon Aug 24 02:53:20 PDT 2026
Author: Pengcheng Wang
Date: 2026-08-24T17:53:15+08:00
New Revision: dfa787605b000d53d0785401ed9bfc66968ad921
URL: https://github.com/llvm/llvm-project/commit/dfa787605b000d53d0785401ed9bfc66968ad921
DIFF: https://github.com/llvm/llvm-project/commit/dfa787605b000d53d0785401ed9bfc66968ad921.diff
LOG: [RISCV] Cost i64 accumulator for Zvdot4a8i partial reductions
Now that `lowerPARTIAL_REDUCE_MLA` can lower an i64 accumulator with
i8 inputs using the dot-product instructions, teach the cost model
to price it so the vectorizer will form it.
- `getPartialReductionCost` accepts an i64 accumulator (reduction
factor 8) and prices the `vdot4a*` plus the reduce-and-accumulate
for both lowering shapes:
- reduce the i32 partial sums with a `vadd.vv` and widen into the
accumulator with a `vwadd.wv` (fixed-length vectors additionally
extract the high i32 subvector with a `vslidedown`), and
- for a single-vector scalable accumulator, whose i32 subvectors are
a fractional LMUL, widen to i64 first and reduce/accumulate with
two i64 `vadd.vv`.
The i64 case requires a wide enough VF (LMUL) to reach the scale-8 factor,
matching how AArch64 only forms it under SVE.
Assisted-by: TRAE CLI (Opus 4.8)
Reviewers: preames, lukel97, topperc
Reviewed By: lukel97, topperc
Pull Request: https://github.com/llvm/llvm-project/pull/216992
Added:
llvm/test/Transforms/LoopVectorize/RISCV/partial-reduce-dot-product-costs.ll
Modified:
llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
Removed:
################################################################################
diff --git a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
index b9d63e9365c75..95d28fd3eaeac 100644
--- a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
@@ -353,17 +353,81 @@ InstructionCost RISCVTTIImpl::getPartialReductionCost(
// zve32x is broken for partial_reduce_umla, but let's make sure we
// don't generate them.
+ // vdot4a* reduces four i8 products into an i32 result; an i64 accumulator is
+ // additionally supported by widening the i32 partial sums to i64 (see
+ // lowerPARTIAL_REDUCE_MLA). VF is the number of i8 input elements, so the
+ // reduction factor is AccumBits / 8 (4 for i32, 8 for i64).
if (!ST->hasStdExtZvdot4a8i() || ST->getELen() < 64 ||
Opcode != Instruction::Add || !BinOp || *BinOp != Instruction::Mul ||
InputTypeA != InputTypeB || !InputTypeA->isIntegerTy(8) ||
- !AccumType->isIntegerTy(32) || !VF.isKnownMultipleOf(4))
+ (!AccumType->isIntegerTy(32) && !AccumType->isIntegerTy(64)))
return InstructionCost::getInvalid();
- Type *Tp = VectorType::get(AccumType, VF.divideCoefficientBy(4));
- std::pair<InstructionCost, MVT> LT = getTypeLegalizationCost(Tp);
+ unsigned Ratio = AccumType->getScalarSizeInBits() / 8;
+ if (!VF.isKnownMultipleOf(Ratio))
+ return InstructionCost::getInvalid();
+
+ // Cost of the vdot4a* itself, which operates on the i32 intermediate type
+ // holding VF/4 elements.
+ Type *DotTp = VectorType::get(Type::getInt32Ty(AccumType->getContext()),
+ VF.divideCoefficientBy(4));
+ std::pair<InstructionCost, MVT> DotLT = getTypeLegalizationCost(DotTp);
// Note: Asuming all vdot4a* variants are equal cost
- return LT.first *
- getRISCVInstructionCost(RISCV::VDOT4A_VV, LT.second, CostKind);
+ InstructionCost Cost =
+ DotLT.first *
+ getRISCVInstructionCost(RISCV::VDOT4A_VV, DotLT.second, CostKind);
+
+ // Account for reducing the i32 partial sums down to the i64 accumulator's
+ // element count and accumulating into it (see lowerPARTIAL_REDUCE_MLA), which
+ // has two shapes depending on the accumulator's LMUL.
+ if (AccumType->isIntegerTy(64)) {
+ LLVMContext &Ctx = AccumType->getContext();
+ Type *I32Ty = Type::getInt32Ty(Ctx);
+ ElementCount AccVF = VF.divideCoefficientBy(Ratio);
+ std::pair<InstructionCost, MVT> AccLT =
+ getTypeLegalizationCost(VectorType::get(AccumType, AccVF));
+
+ // When the i32 subvectors of a single-vector scalable accumulator are a
+ // fractional LMUL, extracting the high subvector would need a vslidedown,
+ // so instead the i32 dot result is widened to i64 first (vsext.vf2 /
+ // vzext.vf2) and then reduced and accumulated with register-aligned i64
+ // vadd.vv.
+ bool WidenFirst = false;
+ if (VF.isScalable() && AccLT.second.isScalableVector()) {
+ MVT NarrowMVT = AccLT.second.changeVectorElementType(MVT::i32);
+ WidenFirst =
+ RISCVVType::decodeVLMUL(RISCVTargetLowering::getLMUL(NarrowMVT))
+ .second;
+ }
+
+ if (WidenFirst) {
+ // The widened i64 dot result has VF/4 elements, i.e. twice the
+ // accumulator's element count, so the reduction plus the accumulate are
+ // two i64 vadd.vv.
+ std::pair<InstructionCost, MVT> WideLT = getTypeLegalizationCost(
+ VectorType::get(AccumType, VF.divideCoefficientBy(4)));
+ Cost +=
+ WideLT.first * getRISCVInstructionCost(RISCV::VSEXT_VF2,
+ WideLT.second, CostKind) +
+ 2 * AccLT.first *
+ getRISCVInstructionCost(RISCV::VADD_VV, AccLT.second, CostKind);
+ } else {
+ // Otherwise the scale-4 i32 sums are halved with a single i32 vadd.vv,
+ // then widened and added into the i64 result with a vwadd.wv.
+ std::pair<InstructionCost, MVT> RedLT =
+ getTypeLegalizationCost(VectorType::get(I32Ty, AccVF));
+ Cost += RedLT.first * getRISCVInstructionCost(RISCV::VADD_VV,
+ RedLT.second, CostKind) +
+ AccLT.first * getRISCVInstructionCost(RISCV::VWADD_WV,
+ AccLT.second, CostKind);
+ // Fixed-length vectors extract the high i32 subvector with a vslidedown.
+ if (VF.isFixed())
+ Cost += DotLT.first * getRISCVInstructionCost(RISCV::VSLIDEDOWN_VI,
+ DotLT.second, CostKind);
+ }
+ }
+
+ return Cost;
}
bool RISCVTTIImpl::shouldExpandReduction(const IntrinsicInst *II) const {
diff --git a/llvm/test/Transforms/LoopVectorize/RISCV/partial-reduce-dot-product-costs.ll b/llvm/test/Transforms/LoopVectorize/RISCV/partial-reduce-dot-product-costs.ll
new file mode 100644
index 0000000000000..fc57c10cc93e4
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/RISCV/partial-reduce-dot-product-costs.ll
@@ -0,0 +1,212 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --filter "Cost.of.*EXPRESSION" --version 6
+; RUN: opt -passes=loop-vectorize -mattr=+v,+experimental-zvdot4a8i \
+; RUN: -tail-folding-policy=dont-fold-tail -debug-only=loop-vectorize \
+; RUN: -disable-output < %s 2>&1 | FileCheck %s
+
+; TODO: Remove -tail-folding-policy=dont-fold-tail when partial reductions with
+; EVL tail folding is supported.
+
+; REQUIRES: asserts
+target triple = "riscv64-none-unknown-elf"
+
+; vdot4a*: i8 inputs into an i32 accumulator (scale-4 partial reduction).
+define i32 @dot_i32(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i32'
+; CHECK: Cost of 1 for VF 8: EXPRESSION vp<[[VP8:%[0-9]+]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i32), (ir<%lb> sext to i32))
+; CHECK: Cost of 1 for VF 8: EXPRESSION vp<[[VP8]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i32), (ir<%lb> sext to i32))
+;
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i32 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = sext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = sext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %add = add i32 %mul, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !0
+exit:
+ ret i32 %add
+}
+
+; i8 inputs into an i64 accumulator, both operands sign-extended (scale-8
+; partial reduction). Costs the vdot4a* plus the extra i32->i64 widen and
+; accumulate; a fixed-length vector additionally needs a vslidedown to extract
+; the high i32 subvector.
+define i64 @dot_i64_ss(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_ss'
+; CHECK: Cost of 4 for VF 8: EXPRESSION vp<[[VP8:%[0-9]+]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i64), (ir<%lb> sext to i64))
+; CHECK: Cost of 4 for VF 8: EXPRESSION vp<[[VP8]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i64), (ir<%lb> sext to i64))
+;
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i64 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = sext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = sext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %conv = sext i32 %mul to i64
+ %add = add i64 %conv, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !0
+exit:
+ ret i64 %add
+}
+
+; Same as above but both operands zero-extended.
+define i64 @dot_i64_uu(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_uu'
+; CHECK: Cost of 4 for VF 8: EXPRESSION vp<[[VP8:%[0-9]+]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> zext to i64), (ir<%lb> zext to i64))
+; CHECK: Cost of 4 for VF 8: EXPRESSION vp<[[VP8]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> zext to i64), (ir<%lb> zext to i64))
+;
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i64 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = zext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = zext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %conv = zext i32 %mul to i64
+ %add = add i64 %conv, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !0
+exit:
+ ret i64 %add
+}
+
+; Mixed sign/zero extension is not a supported dot-product form, so no
+; partial reduction is created for the i64 accumulator.
+define i64 @dot_i64_su(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_su'
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i64 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = sext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = zext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %conv = sext i32 %mul to i64
+ %add = add i64 %conv, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !0
+exit:
+ ret i64 %add
+}
+
+; Scalable single-vector accumulator (nxv1i64): the i32 dot subvectors are a
+; fractional LMUL, so lowerPARTIAL_REDUCE_MLA widens to i64 first (vsext.vf2)
+; and reduces/accumulates with register-aligned i64 vadd.vv instead of a
+; vslidedown + vwadd.wv.
+define i64 @dot_i64_ss_scalable(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_ss_scalable'
+; CHECK: Cost of 5 for VF vscale x 8: EXPRESSION vp<[[VP8:%[0-9]+]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i64), (ir<%lb> sext to i64))
+; CHECK: Cost of 5 for VF vscale x 8: EXPRESSION vp<[[VP8]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i64), (ir<%lb> sext to i64))
+;
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i64 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = sext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = sext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %conv = sext i32 %mul to i64
+ %add = add i64 %conv, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !2
+exit:
+ ret i64 %add
+}
+
+; Same as above but both operands zero-extended.
+define i64 @dot_i64_uu_scalable(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_uu_scalable'
+; CHECK: Cost of 5 for VF vscale x 8: EXPRESSION vp<[[VP8:%[0-9]+]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> zext to i64), (ir<%lb> zext to i64))
+; CHECK: Cost of 5 for VF vscale x 8: EXPRESSION vp<[[VP8]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> zext to i64), (ir<%lb> zext to i64))
+;
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i64 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = zext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = zext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %conv = zext i32 %mul to i64
+ %add = add i64 %conv, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !2
+exit:
+ ret i64 %add
+}
+
+; Wider scalable accumulator (nxv2i64): the i32 dot subvectors are already
+; register-aligned (LMUL 1), so the sums are reduced in i32 (vadd.vv) and
+; accumulated with a single vwadd.wv, matching the fixed-length shape but
+; without the vslidedown.
+define i64 @dot_i64_ss_scalable_m2(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_ss_scalable_m2'
+; CHECK: Cost of 5 for VF vscale x 16: EXPRESSION vp<[[VP8:%[0-9]+]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i64), (ir<%lb> sext to i64))
+; CHECK: Cost of 5 for VF vscale x 16: EXPRESSION vp<[[VP8]]> = ir<%acc> + partial.reduce.add (mul (ir<%la> sext to i64), (ir<%lb> sext to i64))
+;
+entry:
+ br label %loop
+loop:
+ %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+ %acc = phi i64 [ 0, %entry ], [ %add, %loop ]
+ %ga = getelementptr i8, ptr %a, i64 %iv
+ %la = load i8, ptr %ga, align 1
+ %ea = sext i8 %la to i32
+ %gb = getelementptr i8, ptr %b, i64 %iv
+ %lb = load i8, ptr %gb, align 1
+ %eb = sext i8 %lb to i32
+ %mul = mul i32 %ea, %eb
+ %conv = sext i32 %mul to i64
+ %add = add i64 %conv, %acc
+ %iv.next = add i64 %iv, 1
+ %ec = icmp eq i64 %iv.next, 1024
+ br i1 %ec, label %exit, label %loop, !llvm.loop !4
+exit:
+ ret i64 %add
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.vectorize.width", i32 8}
+!2 = distinct !{!2, !1, !3}
+!3 = !{!"llvm.loop.vectorize.scalable.enable"}
+!4 = distinct !{!4, !5, !3}
+!5 = !{!"llvm.loop.vectorize.width", i32 16}
More information about the llvm-commits
mailing list