[RISCV][TTI] Model partial reduce of ext for zvqdotq

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

@@ -303,16 +303,29 @@ InstructionCost RISCVTTIImpl::getPartialReductionCost(
   // zve32x is broken for partial_reduce_umla, but let's make sure we
   // don't generate them.
   if (!ST->hasStdExtZvqdotq() || ST->getELen() < 64 ||
-      Opcode != Instruction::Add || !BinOp || *BinOp != Instruction::Mul ||
-      InputTypeA != InputTypeB || !InputTypeA->isIntegerTy(8) ||
+      Opcode != Instruction::Add || !InputTypeA->isIntegerTy(8) ||
       !AccumType->isIntegerTy(32) || !VF.isKnownMultipleOf(4))
     return InstructionCost::getInvalid();
 
+  // We support both the plain dot product idiom, and the use of dotproduct
+  // to compute a a reduction of an extended value.
+  if (BinOp && (*BinOp != Instruction::Mul || InputTypeA != InputTypeB))
+    return InstructionCost::getInvalid();
+
+  InstructionCost IntMatCost = 0;
+  if (!BinOp) {
+    // Cost to produce one vmv.v.i -- since the constant is shared across any
+    // unrolled copies, don't need to scale by LT.first.
+    Type *Tp = VectorType::get(InputTypeA, VF);
+    std::pair<InstructionCost, MVT> LT = getTypeLegalizationCost(Tp);
+    IntMatCost = getRISCVInstructionCost(RISCV::VMV_V_I, LT.second, CostKind);
+  }
+
   Type *Tp = VectorType::get(AccumType, VF.divideCoefficientBy(4));
   std::pair<InstructionCost, MVT> LT = getTypeLegalizationCost(Tp);
   // Note: Asuming all vqdot* variants are equal cost
-  return LT.first *
-         getRISCVInstructionCost(RISCV::VQDOT_VV, LT.second, CostKind);
+  return IntMatCost + LT.first * getRISCVInstructionCost(RISCV::VQDOT_VV,
+                                                         LT.second, CostKind);
 }
 
 bool RISCVTTIImpl::shouldExpandReduction(const IntrinsicInst *II) const {