[llvm-branch-commits] [RISCV] Cost i64 accumulator for Zvdot4a8i partial reductions (PR #216992)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Aug 18 04:23:48 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Pengcheng Wang (wangpc-pp)
<details>
<summary>Changes</summary>
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), costing the `vdot4a*` plus the extra i32->i64 widen
and accumulate.
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)
---
Full diff: https://github.com/llvm/llvm-project/pull/216992.diff
2 Files Affected:
- (modified) llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp (+26-5)
- (added) llvm/test/Transforms/LoopVectorize/RISCV/partial-reduce-dot-product-costs.ll (+120)
``````````diff
diff --git a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
index 0f636d078324c..772b72708f5dd 100644
--- a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
@@ -353,17 +353,38 @@ 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 the widening extend + accumulate needed for an i64 result.
+ if (AccumType->isIntegerTy(64)) {
+ Type *AccTp = VectorType::get(AccumType, VF.divideCoefficientBy(Ratio));
+ std::pair<InstructionCost, MVT> AccLT = getTypeLegalizationCost(AccTp);
+ Cost += AccLT.first * 2;
+ }
+
+ 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..6a766748522ce
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/RISCV/partial-reduce-dot-product-costs.ll
@@ -0,0 +1,120 @@
+; 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.
+define i64 @dot_i64_ss(ptr %a, ptr %b) {
+; CHECK-LABEL: 'dot_i64_ss'
+; CHECK: Cost of 3 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 3 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 3 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 3 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
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.vectorize.width", i32 8}
``````````
</details>
https://github.com/llvm/llvm-project/pull/216992
More information about the llvm-branch-commits
mailing list