[llvm] [LV] Handle partial sub-reductions with sub in middle block. (PR #178919)

[Florian Hahn via llvm-commits]

================
@@ -5828,7 +5833,46 @@ void VPlanTransforms::createPartialReductions(VPlan &Plan,
     }
   }
 
-  for (const auto &[_, Chains] : ChainsByPhi)
+  for (auto &[Phi, Chains] : ChainsByPhi) {
+    RecurKind RK = cast<VPReductionPHIRecipe>(Phi)->getRecurrenceKind();
+    bool IsTransformed = false;
     for (const VPPartialReductionChain &Chain : Chains)
-      transformToPartialReduction(Chain, Range, CostCtx, Plan);
+      IsTransformed |=
+          transformToPartialReduction(Chain, Range, CostCtx, Plan, RK);
+
+    // Sub-reductions can be implemented in two ways:
+    // (1) negate the operand in the vector loop (the default way).
+    // (2) subtract the reduced value from the init value in the middle block.
+    // Both ways keep the reduction itself as an 'add' reduction.
+    //
+    // The ISD nodes for partial reductions don't support folding the
+    // sub/negation into its operands because the following is not a valid
+    // transformation:
+    //      sub(0, mul(ext(a), ext(b)))
+    //   -> mul(ext(a), ext(sub(0, b)))
+    //
+    // It's therefore better to choose option (2) such that the partial
+    // reduction is always positive (starting at '0') and to do a final
+    // subtract in the middle block.
+    if (IsTransformed && !Range.isEmpty() && RK == RecurKind::Sub) {
+      // Update start value of PHI node.
+      auto *StartInst = cast<VPInstruction>(Phi->getStartValue());
+      assert(StartInst->getOpcode() == VPInstruction::ReductionStartVector);
+      VPValue *OldStartValue = StartInst->getOperand(0);
+      StartInst->setOperand(0, StartInst->getOperand(1));
+
+      // Replace reduction_result by 'sub (startval, reductionresult)'.
+      VPInstruction *RdxResult = vputils::findComputeReductionResult(Phi);
+      assert(RdxResult && "Could not find reduction result");
----------------
fhahn wrote:

Is there a reason we cannot do this directly in `transformToPartialReduction`? I think we already check if it is a last link, could we do it directly there (where we find `ExitValue`?). There's also now VPlan-based sinking, which should take care of sinking outside the loop I think.

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