[llvm] [RISCV][TTI] Model partial reduce of ext for zvqdotq (PR #146788)

[David Sherwood via llvm-commits]

================
@@ -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);
----------------
david-arm wrote:

I'm not really familiar with RISCV lowering, but wouldn't the materialisation in practice be free since only the loop vectoriser uses `getPartialReductionCost` and for vectorised loops the constant should be hoisted out?

https://github.com/llvm/llvm-project/pull/146788