[llvm] [VPlan] Detect and create partial reductions in VPlan. (NFCI) (PR #167851)

[Sam Tebbs via llvm-commits]

================
@@ -991,3 +996,381 @@ bool VPlanTransforms::handleMaxMinNumReductions(VPlan &Plan) {
   MiddleTerm->setOperand(0, NewCond);
   return true;
 }
+
+namespace {
+/// A VPlan-based chain of recipes that form a partial reduction.
+/// Designed to match either:
+///   reduction_bin_op (extend (A), accumulator), or
+///   reduction_bin_op (bin_op (extend (A), (extend (B))), accumulator).
+struct VPPartialReductionChain {
+  VPPartialReductionChain(VPWidenRecipe *Reduction, VPWidenCastRecipe *ExtendA,
+                          VPWidenCastRecipe *ExtendB, VPWidenRecipe *ExtendUser)
+      : Reduction(Reduction), ExtendA(ExtendA), ExtendB(ExtendB),
+        ExtendUser(ExtendUser) {}
+  /// The top-level binary operation that forms the reduction to a scalar
+  /// after the loop body.
+  VPWidenRecipe *Reduction;
+  /// The extension of each of the inner binary operation's operands.
+  VPWidenCastRecipe *ExtendA;
+  VPWidenCastRecipe *ExtendB;
+
+  /// The user of the extends that is then reduced.
+  VPWidenRecipe *ExtendUser;
+};
+
+// Helper to transform a single widen recipe into a partial reduction recipe.
+// Returns true if transformation succeeded.
+static bool transformToPartialReduction(VPWidenRecipe *WidenRecipe,
+                                        unsigned ScaleFactor, VPlan &Plan) {
+  assert(WidenRecipe->getNumOperands() == 2 && "Expected binary operation");
+
+  VPValue *BinOp = WidenRecipe->getOperand(0);
+  VPValue *Accumulator = WidenRecipe->getOperand(1);
+
+  // Swap if needed to ensure Accumulator is the PHI or partial reduction.
+  VPRecipeBase *BinOpRecipe = BinOp->getDefiningRecipe();
+  if (BinOpRecipe && (isa<VPReductionPHIRecipe>(BinOpRecipe) ||
+                      isa<VPPartialReductionRecipe>(BinOpRecipe)))
+    std::swap(BinOp, Accumulator);
+
+  // For chained reductions, only transform if accumulator is already a PHI or
+  // partial reduction. Otherwise, it needs to be transformed first.
+  VPRecipeBase *AccumRecipe = Accumulator->getDefiningRecipe();
+  if (!AccumRecipe || (!isa<VPReductionPHIRecipe>(AccumRecipe) &&
+                       !isa<VPPartialReductionRecipe>(AccumRecipe)))
+    return false;
+
+  if (auto *RdxPhi = dyn_cast<VPReductionPHIRecipe>(AccumRecipe)) {
+    assert(RdxPhi->getVFScaleFactor() == 1 && "scale factor must not be set");
+    RdxPhi->setVFScaleFactor(ScaleFactor);
+
+    // Update ReductionStartVector instruction scale factor.
+    VPValue *StartValue = RdxPhi->getOperand(0);
+    auto *StartInst = cast<VPInstruction>(StartValue);
+    assert(StartInst->getOpcode() == VPInstruction::ReductionStartVector);
+    auto *NewScaleFactor = Plan.getConstantInt(32, ScaleFactor);
+    StartInst->setOperand(2, NewScaleFactor);
+  }
+
+  // Handle SUB by negating the operand and using ADD for the partial reduction.
+  unsigned ReductionOpcode = WidenRecipe->getOpcode();
+  if (ReductionOpcode == Instruction::Sub) {
+    VPBuilder Builder(WidenRecipe);
+
+    // Infer the scalar type for creating the zero constant
+    Type *ElemTy = VPTypeAnalysis(Plan).inferScalarType(BinOp);
+    auto *Zero = Plan.getConstantInt(ElemTy, 0);
+
+    // Create a negation: 0 - BinOp
+    VPIRFlags Flags;
+    if (auto *I = WidenRecipe->getUnderlyingInstr())
+      Flags = VPIRFlags(*I);
+
+    auto *NegRecipe = new VPWidenRecipe(Instruction::Sub, {Zero, BinOp}, Flags,
+                                        VPIRMetadata(), DebugLoc());
+    Builder.insert(NegRecipe);
+    BinOp = NegRecipe;
+    ReductionOpcode = Instruction::Add;
+  }
+
+  VPValue *Cond = nullptr;
+  VPValue *ExitValue = nullptr;
+  if (auto *RedPhiR = dyn_cast<VPReductionPHIRecipe>(Accumulator)) {
+    ExitValue = findComputeReductionResult(RedPhiR)->getOperand(1);
+    match(ExitValue, m_Select(m_VPValue(Cond), m_VPValue(), m_VPValue()));
----------------
SamTebbs33 wrote:

Why do we care about matching a select here? As far as I can see, none of the removed code cares about selects, and the select for the cond is added at execution time, so I'm not sure why we need to care about selects here.

Same for line 1090.

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