Address comments. Add RUN line to test files. Change a comment.

Remove Mask as operand in Partial Reduction Recipe.
Instead just mask the input when creating the recipe.

Author: JamesChesterman

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8898,8 +8898,13 @@ VPRecipeBuilder::tryToCreatePartialReduction(Instruction *Reduction,
     std::swap(BinOp, Accumulator);
 
   VPValue *Mask = getBlockInMask(Reduction->getParent());
-
-  return new VPPartialReductionRecipe(Reduction->getOpcode(), BinOp, Accumulator, Mask, Reduction);
+  if (Mask) {
+    VPValue *Zero =
+        Plan.getOrAddLiveIn(ConstantInt::get(Reduction->getType(), 0));
+    BinOp = Builder.createSelect(Mask, BinOp, Zero, Reduction->getDebugLoc());
+  }
+  return new VPPartialReductionRecipe(Reduction->getOpcode(), BinOp,
+                                      Accumulator, Reduction);
 }
 
 void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
@@ -9711,11 +9716,9 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
     // beginning of the dedicated latch block.
     auto *OrigExitingVPV = PhiR->getBackedgeValue();
     auto *NewExitingVPV = PhiR->getBackedgeValue();
-    // Don't add selects here for partial reductions because the phi and partial
-    // reduction values have less vector elements than Cond. But, each operand
-    // in a select instruction needs to have the same number of vector elements,
-    // so the compiler would crash. Instead, a select, with the active lane
-    // mask, is applied to the inputs to the partial reduction.
+    // Don't output selects for partial reductions because they have an output
+    // with fewer lanes than the VF. So the operands of the select would have
+    // different numbers of lanes. Partial reductions mask the input instead.
     if (!PhiR->isInLoop() && CM.foldTailByMasking() &&
         !isa<VPPartialReductionRecipe>(OrigExitingVPV->getDefiningRecipe())) {
       VPValue *Cond = RecipeBuilder.getBlockInMask(OrigLoop->getHeader());

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -2455,11 +2455,10 @@ class VPPartialReductionRecipe : public VPSingleDefRecipe {
 
 public:
   VPPartialReductionRecipe(Instruction *ReductionInst, VPValue *Op0,
-                           VPValue *Op1, VPValue *Mask = nullptr)
-      : VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1, Mask,
+                           VPValue *Op1)
+      : VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1,
                                  ReductionInst) {}
   VPPartialReductionRecipe(unsigned Opcode, VPValue *Op0, VPValue *Op1,
-                           VPValue *Mask = nullptr,
                            Instruction *ReductionInst = nullptr)
       : VPSingleDefRecipe(VPDef::VPPartialReductionSC,
                           ArrayRef<VPValue *>({Op0, Op1}), ReductionInst),
@@ -2469,18 +2468,12 @@ class VPPartialReductionRecipe : public VPSingleDefRecipe {
     assert((isa<VPReductionPHIRecipe>(AccumulatorRecipe) ||
             isa<VPPartialReductionRecipe>(AccumulatorRecipe)) &&
            "Unexpected operand order for partial reduction recipe");
-    if (Mask)
-      addOperand(Mask);
   }
   ~VPPartialReductionRecipe() override = default;
 
   VPPartialReductionRecipe *clone() override {
     return new VPPartialReductionRecipe(Opcode, getOperand(0), getOperand(1),
-                                        getMask(), getUnderlyingInstr());
-  }
-
-  VPValue *getMask() const {
-    return getNumOperands() == 3 ? getOperand(2) : nullptr;
+                                        getUnderlyingInstr());
   }
 
   VP_CLASSOF_IMPL(VPDef::VPPartialReductionSC)

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -283,6 +283,13 @@ VPPartialReductionRecipe::computeCost(ElementCount VF,
                                       VPCostContext &Ctx) const {
   std::optional<unsigned> Opcode = std::nullopt;
   VPRecipeBase *BinOpR = getOperand(0)->getDefiningRecipe();
+
+  // If the partial reduction is predicated, a select will be operand 0 rather
+  // than the binary op
+  using namespace llvm::VPlanPatternMatch;
+  if (match(getOperand(0), m_Select(m_VPValue(), m_VPValue(), m_VPValue())))
+    BinOpR = BinOpR->getOperand(1)->getDefiningRecipe();
+
   if (auto *WidenR = dyn_cast<VPWidenRecipe>(BinOpR))
     Opcode = std::make_optional(WidenR->getOpcode());
 
@@ -327,12 +334,6 @@ void VPPartialReductionRecipe::execute(VPTransformState &State) {
 
   Type *RetTy = PhiVal->getType();
 
-  VPValue *Mask = getMask();
-  if (Mask) {
-    Value *MaskVal = State.get(Mask);
-    Value *Zero = ConstantInt::get(BinOpVal->getType(), 0);
-    BinOpVal = Builder.CreateSelect(MaskVal, BinOpVal, Zero);
-  }
   CallInst *V = Builder.CreateIntrinsic(
       RetTy, Intrinsic::experimental_vector_partial_reduce_add,
       {PhiVal, BinOpVal}, nullptr, "partial.reduce");