[llvm] [LoopVectorizer] Allow partial reductions to be made in predicated loops (PR #124268)

[Benjamin Maxwell via llvm-commits]

================
@@ -2442,28 +2442,41 @@ class VPReductionPHIRecipe : public VPHeaderPHIRecipe,
 /// A recipe for forming partial reductions. In the loop, an accumulator and
 /// vector operand are added together and passed to the next iteration as the
 /// next accumulator. After the loop body, the accumulator is reduced to a
-/// scalar value.
+/// scalar value. If the mask operand is not nullptr then it is applied to the
+/// vector operand on each iteration.
 class VPPartialReductionRecipe : public VPSingleDefRecipe {
   unsigned Opcode;
 
 public:
   VPPartialReductionRecipe(Instruction *ReductionInst, VPValue *Op0,
-                           VPValue *Op1)
-      : VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1,
+                           VPValue *Op1, VPValue *Mask = nullptr)
+      : VPPartialReductionRecipe(ReductionInst->getOpcode(), Op0, Op1, Mask,
                                  ReductionInst) {}
   VPPartialReductionRecipe(unsigned Opcode, VPValue *Op0, VPValue *Op1,
+                           VPValue *Mask = nullptr,
                            Instruction *ReductionInst = nullptr)
       : VPSingleDefRecipe(VPDef::VPPartialReductionSC,
                           ArrayRef<VPValue *>({Op0, Op1}), ReductionInst),
         Opcode(Opcode) {
     assert(isa<VPReductionPHIRecipe>(getOperand(1)->getDefiningRecipe()) &&
            "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),
-                                        getUnderlyingInstr());
+    return getNumOperands() == 3
+               ? new VPPartialReductionRecipe(Opcode, getOperand(0),
+                                              getOperand(1), getOperand(2),
+                                              getUnderlyingInstr())
+               : new VPPartialReductionRecipe(Opcode, getOperand(0),
+                                              getOperand(1), nullptr,
+                                              getUnderlyingInstr());
----------------
MacDue wrote:

```suggestion
    return  new VPPartialReductionRecipe(Opcode, getOperand(0),
                                              getOperand(1), getMask()
                                              getUnderlyingInstr());
```

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