[VPlan] Introduce ComputeReductionResult VPInstruction opcode.

[fhahn]

This patch introduces a new ComputeReductionResult opcode to compute the final reduction result in the middle block. The code from fixReduction has been moved to ComputeReductionResult, after some earlier cleanup changes to model parts of fixReduction explicitly elsewhere as needed.

The recipe may be broken down further in the future.

Note that the ComputeReductionResult at the moment also creates the phi node to merge the reduction result from the trip count check and the middle block, to be used as resume value for the scalar remainder loop. Once we have a VPValue for the reduction result, this can also be modeled explicitly and moved out of the recipe.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

Changes

This patch introduces a new ComputeReductionResult opcode to compute the final reduction result in the middle block. The code from fixReduction has been moved to ComputeReductionResult, after some earlier cleanup changes to model parts of fixReduction explicitly elsewhere as needed.

The recipe may be broken down further in the future.

Note that the ComputeReductionResult at the moment also creates the phi node to merge the reduction result from the trip count check and the middle block, to be used as resume value for the scalar remainder loop. Once we have a VPValue for the reduction result, this can also be modeled explicitly and moved out of the recipe.

---

Patch is 28.81 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/70253.diff

7 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h (+7-4)
• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+64-190)
• (modified) llvm/lib/Transforms/Vectorize/VPlan.cpp (+1)
• (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+2-1)
• (modified) llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp (+136)
• (modified) llvm/test/Transforms/LoopVectorize/first-order-recurrence-sink-replicate-region.ll (+2-1)
• (modified) llvm/test/Transforms/LoopVectorize/vplan-printing.ll (+5-2)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
index 577ce8000de27b9..a4ead0523a2fb70 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
@@ -347,10 +347,13 @@ class LoopVectorizationPlanner {
   /// vectorization re-using plans for both the main and epilogue vector loops.
   /// It should be removed once the re-use issue has been fixed.
   /// \p ExpandedSCEVs is passed during execution of the plan for epilogue loop
-  /// to re-use expansion results generated during main plan execution. Returns
-  /// a mapping of SCEVs to their expanded IR values. Note that this is a
-  /// temporary workaround needed due to the current epilogue handling.
-  DenseMap<const SCEV *, Value *>
+  /// to re-use expansion results generated during main plan execution.
+  ///
+  /// Returns a mapping of SCEVs to their expanded IR values and a mapping for
+  /// the reduction resume values. Note that this is a temporary workaround
+  /// needed due to the current epilogue handling.
+  std::pair<DenseMap<const SCEV *, Value *>,
+            DenseMap<const RecurrenceDescriptor *, Value *>>
   executePlan(ElementCount VF, unsigned UF, VPlan &BestPlan,
               InnerLoopVectorizer &LB, DominatorTree *DT,
               bool IsEpilogueVectorization,
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 0807d2a7e5a2671..13267ec5017c181 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -566,10 +566,6 @@ class InnerLoopVectorizer {
   /// able to vectorize with strict in-order reductions for the given RdxDesc.
   bool useOrderedReductions(const RecurrenceDescriptor &RdxDesc);
 
-  // Returns the resume value (bc.merge.rdx) for a reduction as
-  // generated by fixReduction.
-  PHINode *getReductionResumeValue(const RecurrenceDescriptor &RdxDesc);
-
   /// Create a new phi node for the induction variable \p OrigPhi to resume
   /// iteration count in the scalar epilogue, from where the vectorized loop
   /// left off. \p Step is the SCEV-expanded induction step to use. In cases
@@ -1151,14 +1147,6 @@ void InnerLoopVectorizer::collectPoisonGeneratingRecipes(
   }
 }
 
-PHINode *InnerLoopVectorizer::getReductionResumeValue(
-    const RecurrenceDescriptor &RdxDesc) {
-  auto It = ReductionResumeValues.find(&RdxDesc);
-  assert(It != ReductionResumeValues.end() &&
-         "Expected to find a resume value for the reduction.");
-  return It->second;
-}
-
 namespace llvm {
 
 // Loop vectorization cost-model hints how the scalar epilogue loop should be
@@ -3617,11 +3605,6 @@ void InnerLoopVectorizer::fixCrossIterationPHIs(VPTransformState &State) {
   VPBasicBlock *Header =
       State.Plan->getVectorLoopRegion()->getEntryBasicBlock();
 
-  for (VPRecipeBase &R : Header->phis()) {
-    if (auto *ReductionPhi = dyn_cast<VPReductionPHIRecipe>(&R))
-      fixReduction(ReductionPhi, State);
-  }
-
   for (VPRecipeBase &R : Header->phis()) {
     if (auto *FOR = dyn_cast<VPFirstOrderRecurrencePHIRecipe>(&R))
       fixFixedOrderRecurrence(FOR, State);
@@ -3747,169 +3730,6 @@ void InnerLoopVectorizer::fixFixedOrderRecurrence(
   Phi->setName("scalar.recur");
 }
 
-void InnerLoopVectorizer::fixReduction(VPReductionPHIRecipe *PhiR,
-                                       VPTransformState &State) {
-  PHINode *OrigPhi = cast<PHINode>(PhiR->getUnderlyingValue());
-  // Get it's reduction variable descriptor.
-  assert(Legal->isReductionVariable(OrigPhi) &&
-         "Unable to find the reduction variable");
-  const RecurrenceDescriptor &RdxDesc = PhiR->getRecurrenceDescriptor();
-
-  RecurKind RK = RdxDesc.getRecurrenceKind();
-  TrackingVH<Value> ReductionStartValue = RdxDesc.getRecurrenceStartValue();
-  Instruction *LoopExitInst = RdxDesc.getLoopExitInstr();
-  if (auto *I = dyn_cast<Instruction>(&*ReductionStartValue))
-    State.setDebugLocFrom(I->getDebugLoc());
-
-  VPValue *LoopExitInstDef = PhiR->getBackedgeValue();
-
-  // Before each round, move the insertion point right between
-  // the PHIs and the values we are going to write.
-  // This allows us to write both PHINodes and the extractelement
-  // instructions.
-  Builder.SetInsertPoint(LoopMiddleBlock,
-                         LoopMiddleBlock->getFirstInsertionPt());
-
-  State.setDebugLocFrom(LoopExitInst->getDebugLoc());
-
-  Type *PhiTy = OrigPhi->getType();
-  // If tail is folded by masking, the vector value to leave the loop should be
-  // a Select choosing between the vectorized LoopExitInst and vectorized Phi,
-  // instead of the former. For an inloop reduction the reduction will already
-  // be predicated, and does not need to be handled here.
-  if (Cost->foldTailByMasking() && !PhiR->isInLoop()) {
-    VPValue *Def = nullptr;
-    for (VPUser *U : LoopExitInstDef->users()) {
-      auto *S = dyn_cast<VPInstruction>(U);
-      if (S && S->getOpcode() == Instruction::Select) {
-        Def = S;
-        break;
-      }
-    }
-    if (Def)
-      LoopExitInstDef = Def;
-  }
-
-  VectorParts RdxParts(UF);
-  for (unsigned Part = 0; Part < UF; ++Part)
-    RdxParts[Part] = State.get(LoopExitInstDef, Part);
-
-  // If the vector reduction can be performed in a smaller type, we truncate
-  // then extend the loop exit value to enable InstCombine to evaluate the
-  // entire expression in the smaller type.
-  if (VF.isVector() && PhiTy != RdxDesc.getRecurrenceType()) {
-    Builder.SetInsertPoint(LoopMiddleBlock,
-                           LoopMiddleBlock->getFirstInsertionPt());
-    Type *RdxVecTy = VectorType::get(RdxDesc.getRecurrenceType(), VF);
-    for (unsigned Part = 0; Part < UF; ++Part) {
-      RdxParts[Part] = Builder.CreateTrunc(RdxParts[Part], RdxVecTy);
-    }
-  }
-
-  // Reduce all of the unrolled parts into a single vector.
-  Value *ReducedPartRdx = RdxParts[0];
-  unsigned Op = RecurrenceDescriptor::getOpcode(RK);
-
-  // The middle block terminator has already been assigned a DebugLoc here (the
-  // OrigLoop's single latch terminator). We want the whole middle block to
-  // appear to execute on this line because: (a) it is all compiler generated,
-  // (b) these instructions are always executed after evaluating the latch
-  // conditional branch, and (c) other passes may add new predecessors which
-  // terminate on this line. This is the easiest way to ensure we don't
-  // accidentally cause an extra step back into the loop while debugging.
-  State.setDebugLocFrom(LoopMiddleBlock->getTerminator()->getDebugLoc());
-  if (PhiR->isOrdered())
-    ReducedPartRdx = RdxParts[UF - 1];
-  else {
-    // Floating-point operations should have some FMF to enable the reduction.
-    IRBuilderBase::FastMathFlagGuard FMFG(Builder);
-    Builder.setFastMathFlags(RdxDesc.getFastMathFlags());
-    for (unsigned Part = 1; Part < UF; ++Part) {
-      Value *RdxPart = RdxParts[Part];
-      if (Op != Instruction::ICmp && Op != Instruction::FCmp)
-        ReducedPartRdx = Builder.CreateBinOp(
-            (Instruction::BinaryOps)Op, RdxPart, ReducedPartRdx, "bin.rdx");
-      else if (RecurrenceDescriptor::isAnyOfRecurrenceKind(RK))
-        ReducedPartRdx = createAnyOfOp(Builder, ReductionStartValue, RK,
-                                       ReducedPartRdx, RdxPart);
-      else
-        ReducedPartRdx = createMinMaxOp(Builder, RK, ReducedPartRdx, RdxPart);
-    }
-  }
-
-  // Create the reduction after the loop. Note that inloop reductions create the
-  // target reduction in the loop using a Reduction recipe.
-  if (VF.isVector() && !PhiR->isInLoop()) {
-    ReducedPartRdx =
-        createTargetReduction(Builder, RdxDesc, ReducedPartRdx, OrigPhi);
-    // If the reduction can be performed in a smaller type, we need to extend
-    // the reduction to the wider type before we branch to the original loop.
-    if (PhiTy != RdxDesc.getRecurrenceType())
-      ReducedPartRdx = RdxDesc.isSigned()
-                           ? Builder.CreateSExt(ReducedPartRdx, PhiTy)
-                           : Builder.CreateZExt(ReducedPartRdx, PhiTy);
-  }
-
-  PHINode *ResumePhi =
-      dyn_cast<PHINode>(PhiR->getStartValue()->getUnderlyingValue());
-
-  // Create a phi node that merges control-flow from the backedge-taken check
-  // block and the middle block.
-  PHINode *BCBlockPhi = PHINode::Create(PhiTy, 2, "bc.merge.rdx",
-                                        LoopScalarPreHeader->getTerminator());
-
-  // If we are fixing reductions in the epilogue loop then we should already
-  // have created a bc.merge.rdx Phi after the main vector body. Ensure that
-  // we carry over the incoming values correctly.
-  for (auto *Incoming : predecessors(LoopScalarPreHeader)) {
-    if (Incoming == LoopMiddleBlock)
-      BCBlockPhi->addIncoming(ReducedPartRdx, Incoming);
-    else if (ResumePhi && llvm::is_contained(ResumePhi->blocks(), Incoming))
-      BCBlockPhi->addIncoming(ResumePhi->getIncomingValueForBlock(Incoming),
-                              Incoming);
-    else
-      BCBlockPhi->addIncoming(ReductionStartValue, Incoming);
-  }
-
-  // Set the resume value for this reduction
-  ReductionResumeValues.insert({&RdxDesc, BCBlockPhi});
-
-  // If there were stores of the reduction value to a uniform memory address
-  // inside the loop, create the final store here.
-  if (StoreInst *SI = RdxDesc.IntermediateStore) {
-    StoreInst *NewSI =
-        Builder.CreateAlignedStore(ReducedPartRdx, SI->getPointerOperand(),
-                                   SI->getAlign());
-    propagateMetadata(NewSI, SI);
-
-    // If the reduction value is used in other places,
-    // then let the code below create PHI's for that.
-  }
-
-  // Now, we need to fix the users of the reduction variable
-  // inside and outside of the scalar remainder loop.
-
-  // We know that the loop is in LCSSA form. We need to update the PHI nodes
-  // in the exit blocks.  See comment on analogous loop in
-  // fixFixedOrderRecurrence for a more complete explaination of the logic.
-  if (!Cost->requiresScalarEpilogue(VF.isVector()))
-    for (PHINode &LCSSAPhi : LoopExitBlock->phis())
-      if (llvm::is_contained(LCSSAPhi.incoming_values(), LoopExitInst)) {
-        LCSSAPhi.addIncoming(ReducedPartRdx, LoopMiddleBlock);
-        State.Plan->removeLiveOut(&LCSSAPhi);
-      }
-
-  // Fix the scalar loop reduction variable with the incoming reduction sum
-  // from the vector body and from the backedge value.
-  int IncomingEdgeBlockIdx =
-      OrigPhi->getBasicBlockIndex(OrigLoop->getLoopLatch());
-  assert(IncomingEdgeBlockIdx >= 0 && "Invalid block index");
-  // Pick the other block.
-  int SelfEdgeBlockIdx = (IncomingEdgeBlockIdx ? 0 : 1);
-  OrigPhi->setIncomingValue(SelfEdgeBlockIdx, BCBlockPhi);
-  OrigPhi->setIncomingValue(IncomingEdgeBlockIdx, LoopExitInst);
-}
-
 void InnerLoopVectorizer::sinkScalarOperands(Instruction *PredInst) {
   // The basic block and loop containing the predicated instruction.
   auto *PredBB = PredInst->getParent();
@@ -7641,7 +7461,9 @@ static void AddRuntimeUnrollDisableMetaData(Loop *L) {
   }
 }
 
-SCEV2ValueTy LoopVectorizationPlanner::executePlan(
+std::pair<DenseMap<const SCEV *, Value *>,
+          DenseMap<const RecurrenceDescriptor *, Value *>>
+LoopVectorizationPlanner::executePlan(
     ElementCount BestVF, unsigned BestUF, VPlan &BestVPlan,
     InnerLoopVectorizer &ILV, DominatorTree *DT, bool IsEpilogueVectorization,
     const DenseMap<const SCEV *, Value *> *ExpandedSCEVs) {
@@ -7669,6 +7491,7 @@ SCEV2ValueTy LoopVectorizationPlanner::executePlan(
     State.Builder.SetInsertPoint(OrigLoop->getLoopPreheader()->getTerminator());
     BestVPlan.getPreheader()->execute(&State);
   }
+  State.CFG.VPBB2IRBB[BestVPlan.getPreheader()] = OrigLoop->getLoopPreheader();
   if (!ILV.getTripCount())
     ILV.setTripCount(State.get(BestVPlan.getTripCount(), {0, 0}));
   else
@@ -7719,6 +7542,27 @@ SCEV2ValueTy LoopVectorizationPlanner::executePlan(
 
   BestVPlan.execute(&State);
 
+  DenseMap<const RecurrenceDescriptor *, Value *> ReductionResumeValues;
+  for (VPRecipeBase &R : *cast<VPBasicBlock>(
+           BestVPlan.getVectorLoopRegion()->getSingleSuccessor())) {
+
+    auto *Red = dyn_cast<VPInstruction>(&R);
+    if (!Red || Red->getOpcode() != VPInstruction::ComputeReductionResult)
+      continue;
+
+    Value *FinalValue = State.get(
+        Red, VPIteration(State.UF - 1, VPLane::getLastLaneForVF(State.VF)));
+    auto *PhiR = cast<VPReductionPHIRecipe>(Red->getOperand(0));
+    assert(FinalValue->hasOneUse() || PhiR->isInLoop() ||
+           PhiR->getRecurrenceDescriptor().IntermediateStore);
+
+    PHINode *MergePhi = nullptr;
+    for (Value *U : FinalValue->users())
+      if (auto *P = dyn_cast<PHINode>(U))
+        MergePhi = P;
+    ReductionResumeValues[&PhiR->getRecurrenceDescriptor()] = MergePhi;
+  }
+
   // Keep all loop hints from the original loop on the vector loop (we'll
   // replace the vectorizer-specific hints below).
   MDNode *OrigLoopID = OrigLoop->getLoopID();
@@ -7751,8 +7595,7 @@ SCEV2ValueTy LoopVectorizationPlanner::executePlan(
   ILV.fixVectorizedLoop(State, BestVPlan);
 
   ILV.printDebugTracesAtEnd();
-
-  return State.ExpandedSCEVs;
+  return {State.ExpandedSCEVs, ReductionResumeValues};
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
@@ -9054,7 +8897,25 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
   SmallVector<VPReductionPHIRecipe *> InLoopReductionPhis;
   for (VPRecipeBase &R : Header->phis()) {
     auto *PhiR = dyn_cast<VPReductionPHIRecipe>(&R);
-    if (!PhiR || !PhiR->isInLoop() || (MinVF.isScalar() && !PhiR->isOrdered()))
+    if (!PhiR)
+      continue;
+
+    // TODO: At the moment ComputeReductionResult also creates the bc.merge.rdx
+    // phi nodes, hence it needs to be created unconditionally here, until the
+    // reduction resume value handling is also modeled in VPlan.
+    VPInstruction *FinalReductionResult =
+        new VPInstruction(VPInstruction::ComputeReductionResult,
+                          {PhiR, PhiR->getBackedgeValue()});
+    cast<VPBasicBlock>(Plan->getVectorLoopRegion()->getSingleSuccessor())
+        ->appendRecipe(FinalReductionResult);
+    for (VPUser *U : to_vector(PhiR->getBackedgeValue()->users())) {
+      auto *LiveOut = dyn_cast<VPLiveOut>(U);
+      if (!LiveOut)
+        continue;
+      LiveOut->setOperand(0, FinalReductionResult);
+    }
+
+    if (!PhiR->isInLoop() || (MinVF.isScalar() && !PhiR->isOrdered()))
       continue;
     InLoopReductionPhis.push_back(PhiR);
   }
@@ -9089,9 +8950,13 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
     for (VPRecipeBase *CurrentLink : Worklist.getArrayRef().drop_front()) {
       VPValue *PreviousLinkV = PreviousLink->getVPSingleValue();
 
+      if (CurrentLink->getParent()->getParent() !=
+          PhiR->getParent()->getParent()) {
+        continue;
+      }
       Instruction *CurrentLinkI = CurrentLink->getUnderlyingInstr();
 
-      // Index of the first operand which holds a non-mask vector operand.
+      // Index of Rhe first operand which holds a non-mask vector operand.
       unsigned IndexOfFirstOperand;
       // Recognize a call to the llvm.fmuladd intrinsic.
       bool IsFMulAdd = (Kind == RecurKind::FMulAdd);
@@ -9100,7 +8965,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
       if (IsFMulAdd) {
         assert(
             RecurrenceDescriptor::isFMulAddIntrinsic(CurrentLinkI) &&
-            "Expected instruction to be a call to the llvm.fmuladd intrinsic");
+            "ExpectedRinstruction to be a call to the llvm.fmuladd intrinsic");
         assert(((MinVF.isScalar() && isa<VPReplicateRecipe>(CurrentLink)) ||
                 isa<VPWidenCallRecipe>(CurrentLink)) &&
                CurrentLink->getOperand(2) == PreviousLinkV &&
@@ -9189,6 +9054,14 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
               ? new VPInstruction(Instruction::Select, {Cond, Red, PhiR}, FMFs)
               : new VPInstruction(Instruction::Select, {Cond, Red, PhiR});
       Result->insertBefore(&*Builder.getInsertPoint());
+      for (VPUser *U : Red->users()) {
+        auto *UR = cast<VPRecipeBase>(U);
+        if (UR->getParent()->getParent() == PhiR->getParent()->getParent())
+          continue;
+        assert(cast<VPInstruction>(UR)->getOpcode() ==
+               VPInstruction::ComputeReductionResult);
+        UR->setOperand(1, Result->getVPSingleValue());
+      }
       if (PreferPredicatedReductionSelect ||
           TTI.preferPredicatedReductionSelect(
               PhiR->getRecurrenceDescriptor().getOpcode(), PhiTy,
@@ -10197,8 +10070,8 @@ bool LoopVectorizePass::processLoop(Loop *L) {
                                            EPI, &LVL, &CM, BFI, PSI, Checks);
 
         VPlan &BestMainPlan = LVP.getBestPlanFor(EPI.MainLoopVF);
-        auto ExpandedSCEVs = LVP.executePlan(EPI.MainLoopVF, EPI.MainLoopUF,
-                                             BestMainPlan, MainILV, DT, true);
+        const auto &[ExpandedSCEVs, ReductionResumeValues] = LVP.executePlan(
+            EPI.MainLoopVF, EPI.MainLoopUF, BestMainPlan, MainILV, DT, true);
         ++LoopsVectorized;
 
         // Second pass vectorizes the epilogue and adjusts the control flow
@@ -10239,8 +10112,9 @@ bool LoopVectorizePass::processLoop(Loop *L) {
           Value *ResumeV = nullptr;
           // TODO: Move setting of resume values to prepareToExecute.
           if (auto *ReductionPhi = dyn_cast<VPReductionPHIRecipe>(&R)) {
-            ResumeV = MainILV.getReductionResumeValue(
-                ReductionPhi->getRecurrenceDescriptor());
+            ResumeV = ReductionResumeValues
+                          .find(&ReductionPhi->getRecurrenceDescriptor())
+                          ->second;
           } else {
             // Create induction resume values for both widened pointer and
             // integer/fp inductions and update the start value of the induction
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 583f652eed104e0..ecb5bd5366b9ed7 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -456,6 +456,7 @@ void VPBasicBlock::execute(VPTransformState *State) {
     // The Exit block of a loop is always set to be successor 0 of the Exiting
     // block.
     cast<BranchInst>(ExitingBB->getTerminator())->setSuccessor(0, NewBB);
+    State->Builder.SetInsertPoint(NewBB->getFirstNonPHI());
   } else if (PrevVPBB && /* A */
              !((SingleHPred = getSingleHierarchicalPredecessor()) &&
                SingleHPred->getExitingBasicBlock() == PrevVPBB &&
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index e65a7ab2cd028ee..614d78b73e7329d 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1034,7 +1034,8 @@ class VPInstruction : public VPRecipeWithIRFlags, public VPValue {
     // canonical IV separately for each unrolled part.
     CanonicalIVIncrementForPart,
     BranchOnCount,
-    BranchOnCond
+    BranchOnCond,
+    ComputeReductionResult,
   };
 
 private:
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index efc95c1cd08c6fd..bd156c73f456b22 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -27,6 +27,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include <cassert>
 
@@ -403,6 +404,138 @@ Value *VPInstruction::generateInstruction(VPTransformState &State,
     Builder.GetInsertBlock()->getTerminator()->eraseFromParent();
     return CondBr;
   }
+  case VPInstruction::ComputeReductionResult: {
+    if (Part != 0)
+      return State.get(
+          this, VPIteration(State.UF - 1, VPLane::getLastLaneForVF(State.VF)));
+
+    auto *PhiR = dyn_cast<VPReductionPHIRecipe>(getOperand(0));
+    PHINode *OrigPhi = cast<PHINode>(PhiR->getUnderlyingValue());
+    // Get it's reduction variable descriptor.
+    const RecurrenceDescriptor &RdxDesc = PhiR...
[truncated]

[alexey-bataev]

break; ? to early exit out of the loop if the first phi is found?

[fhahn]

Adjusted thanks! There actually can be multiple phi users (one could be the phi node in the vector loop), refined to check that the phi is in the original scalar PH.

[alexey-bataev]

Suggested change

	VPInstruction *FinalReductionResult =
	auto *FinalReductionResult =

[fhahn]

Done, thanks!

[alexey-bataev]

Misprint?

[fhahn]

Fixed, thanks!

[alexey-bataev]

Misprint?

[fhahn]

undone, thanks!

[alexey-bataev]

Hm, a reference to a copied containers? May it cause troubles use-after-delete? Should it be rvalue ref instead?

[fhahn]

I think the const reference should extend the lifetime of the returned value.

[alexey-bataev]

Suggested change

	PHINode *ResumePhi =
	auto *ResumePhi =

[alexey-bataev]

Suggested change

	auto *OrigLoop = State.LI->getLoopFor(OrigPhi->getParent());
	const Loop *OrigLoop = State.LI->getLoopFor(OrigPhi->getParent());

[fhahn]

Done, thanks!

[ayalz]

nit: not sure I see it right - should this be auto or const Loop now?

[fhahn]

Changed to const Loop (the type isn't entirely obvious like with dyn_cast)

[alexey-bataev]

Suggested change

	PHINode *BCBlockPhi = PHINode::Create(PhiTy, 2, "bc.merge.rdx",
	auto *BCBlockPhi = PHINode::Create(PhiTy, 2, "bc.merge.rdx",

[fhahn]

Done, thanks!

[alexey-bataev]

Suggested change

	else if (ResumePhi && llvm::is_contained(ResumePhi->blocks(), Incoming))
	else if (ResumePhi && is_contained(ResumePhi->blocks(), Incoming))

[fhahn]

Done, thanks!

[alexey-bataev]

Suggested change

	StoreInst *NewSI = Builder.CreateAlignedStore(
	auto *NewSI = Builder.CreateAlignedStore(

[fhahn]

Done, thanks!

[fhahn]

Address comments & ping :)

Apologies if the force-push messed anything up

[fhahn]

Adjusted thanks! There actually can be multiple phi users (one could be the phi node in the vector loop), refined to check that the phi is in the original scalar PH.

[fhahn]

Done, thanks!

[fhahn]

Fixed, thanks!

[fhahn]

undone, thanks!

[fhahn]

Adjusted, thanks!

[fhahn]

Done, thanks!

[fhahn]

Done, thanks!

[fhahn]

Done, thanks!

[fhahn]

Done, thanks!

[fhahn]

I think the const reference should extend the lifetime of the returned value.

[fhahn]

ping :)

[alexey-bataev]

Better to outline this as lambda

[fhahn]

Outlined, thanks!

[alexey-bataev]

Drop braces

[fhahn]

Removed, thanks!

[alexey-bataev]

Add assertion message

[fhahn]

Done, thanks!

[fhahn]

Rebased and addressed comments, thanks!

[fhahn]

Outlined, thanks!

[fhahn]

Removed, thanks!

[fhahn]

Done, thanks!

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[alexey-bataev]

auto It = find_if(FinalValue->users(), [](Value U) {
auto *P = dyn_cast(U);
return P && OrigScalarPH == P->getParent();
});
assert(It != FinalValue->users().end(), "must have merge phi");
return {&PhiR->getRecurrenceDescriptor(), *It};

[fhahn]

Ah forgot about this one, should be fixed now!

[ayalz]

Good step forward!

[ayalz]

nit: above "Generate the IR code for the body of the vectorized loop" is no longer accurate. Suggest to update with something like "Generate the IR code for the vectorized loop captured in VPlan"

[fhahn]

Adjusted, thanks!

[ayalz]

Still awaits adjustment?

[fhahn]

Should be adjusted now, thanks!

[ayalz]

So now better rename fixCrossIterationPHIs() to fixFixedOrderRecurrences(), along with an explanation that inductions and reductions are fixed by VPlan execute? Or leave a TODO for now, until FORs are also fixed by VPlan and this method can be retired.

[fhahn]

Inlined the loop to the caller

[ayalz]

Is setting VPBB2IRBB for the preheader independent of this patch?
Worth updating inside the if, for non-empty preheaders only?
Note that multiple VPBB's may be associated with OrigLoop's preheader.

[fhahn]

Not needed for this patch, removed, thanks!

[ayalz]

Worth commenting, something like
// 2.5 Collect reduction resume values.
?

Worth outlining?

[fhahn]

Outlined and added comment, thanks!

[ayalz]

Suggested change

	// Check if ResResult is a CompueReductionResult instruction, and if it is
	// Check if RedResult is a CompueReductionResult instruction, and if it is

[fhahn]

Fixed, thanks!

[ayalz]

Suggesting a ComputeReductionResult recipe could hold its RdxDesc.

[fhahn]

Would be an option, see my previous comment. There are multiple properties needed from PhiR here as well

[ayalz]

-- which exacerbates the undesired issue of implicit cross-recipe dependences. Please add a FIXME to emphasize that this is an intermediate situation to support gradual refactoring, possibly also clarifying at the summary that, rather than The recipe may be broken down further in the future - it must be.

[fhahn]

Added a FIXME, thanks!

[ayalz]

Isn't ComputeReductionResult placed in LoopMiddleBlock, with the Builder already pointing to it?

[fhahn]

Removed, thanks!

[ayalz]

nit: more accurately called LoopExitingInstDef, but better retain current name in this patch.

[fhahn]

Adjusted and also dropped the Inst part, as this doesn't refer the an instruction.

[ayalz]

commented-out code?

[fhahn]

Not needed, removed, thanks!

[ayalz]

Unrelated to this patch: better handle type size reductions as part of truncateToMinBW, along with all other type shrinkings, rather than here.

[fhahn]

Agreed, added a TODO

[ayalz]

Good step forward!

[ayalz]

VPLiveOut::fixPhi() indeed looks for the last Part, but it does check if to look for the first or last Lane, according to vputils::isUniformAfterVectorization(ExitValue)?

It's somewhat confusing to set a value for the last Part, when operating upon the first Part. Perhaps more consistent for now to store a value for all parts, as do, e.g., VPExpandSCEVRecipe::execute() and VPCanonicalIVPHIRecipe::execute()?

BTW, VPFirstOrderRecurrencePHIRecipe::execute() sets only the first Part, perhaps it should do so under an if Part == 0? (Same for VPWidenPHIRecipe::execute(), but that applies only to native, i.e., to UF=1)

[ayalz]

-- which highlights the fact that this patch is the first to introduce a recipe in the Middle block -- something worth noting in the summary, if not in the title...

nit: perhaps more consistent to set the insertion point earlier, next to setting CFG.PrevBB, as part of // ExitBB can be re-used ... (or rather it must be reused), rather than under the last // Update the branch ... part. Also, doing so aligns with the else if case below (although arguably setting the insertion point to Terminator is redundant after having set it to NewBB and created Terminator?)

[ayalz]

-- which exacerbates the undesired issue of implicit cross-recipe dependences. Please add a FIXME to emphasize that this is an intermediate situation to support gradual refactoring, possibly also clarifying at the summary that, rather than The recipe may be broken down further in the future - it must be.

[ayalz]

nit: not sure I see it right - should this be auto or const Loop now?

[ayalz]

Perhaps there's a better way to collect these merge phi's per RdxDesc than to lookup ComputeReductionResult's first operand along with first user in OrigScalarPH, see below.

[fhahn]

This isn't needed any longer; the merge phis are created here now, thanks!

[ayalz]

Suggested change

	static void getMergePhiForReduction(
	static void collectMergePhiForReduction(

[fhahn]

Renamed, thanks!

[ayalz]

Ah, sure, same as ExpandedSCEVs.find(ExpandR->getSCEV())->second);

[ayalz]

Is this setting of insertion point to Middle block redundant?

[fhahn]

Yes, removed!

[ayalz]

Should ComputeReductionResult go ahead and record BCBlockPhi inside some State.ReductionResumeValues[&RdxDesc] = BCBlockPhi; when creating it here, similar to how VPExpandSCEVRecipe::execute() records State.ExpandedSCEVs[Expr] = Res;, rather than having getMergePhiForReduction() search for it later. The latter is arguably more encapsulated, the former seems more consistent?

[fhahn]

Hmm, I think it is preferable to keep things encapsulated if possible. I think the main reason for recording the expanded SCEVs in State.ExpandedSCEVs is to avoid duplicated SCEV expansion during epilogue vectorization, due to executing 2 different, unconnected VPlans.

[ayalz]

Encapsulated approach is preferred?

[fhahn]

Yep, missed to submit some outstanding comments, should be visible now.

[fhahn]

Note that the order of the reduction result computation changed in some tests, as now we create them in the original order the reduction recipes have been created

[ayalz]

Additional thoughts for potentially more gradual/consistent refactoring.

[ayalz]

Suggested change

	auto *PhiR = dyn_cast<VPReductionPHIRecipe>(getOperand(0));
	auto *PhiR = cast<VPReductionPHIRecipe>(getOperand(0));

[fhahn]

Done, thanks!

[ayalz]

Still awaits adjustment?

[ayalz]

Documentation should expand, e.g., to explain something like:
// A ComputeReductionResult recipe is added to the middle block, also for in-loop reductions which conmpute their result in-loop, because this recipe currently takes care of generating the subsequent bc.merge.rdx phi which applies to all reductions.

[fhahn]

Expanded thanks, with adjustments w.r.t. bc.merge.rdx phi creation.

[ayalz]

Would it be better to have collectMergePhiForReduction(), in addition to collecting the bc.merge.rdx phi, also take care of generating it inside the scalar preheader - a block that is yet to be modeled in VPlan? Thereby letting the new ComputeReductionResult recipe, placed inside the middle block, take care of generating code there only?

[fhahn]

Moved the code here, it works well, thanks!

[ayalz]

Perhaps better keep the creation of bc.merge.rdx outside VPlan, until it can be moved into a suitable recipe and VPBB?
May still "be created unconditionally": ComputeReductionResult finalizes the computation of non in-loop reductions. For in-loop reductions one could (temporarily?) use ComputeReductionResult to represent (and easily retrieve) the ReductionResult, even if no additional Compute is needed, as a form of LCSSA Phi?

Should LV generate an actual LCSSA Phi, i.e., maintain LCSSA form? Note that all VPInstructions are currently excluded from VPRecipeBase::isPhi().

How do we prevent ComputeReductionResult from being treated as dead code, if-not/until hooked-up to a LiveOut.

[fhahn]

Creation of bc.merge.rdx has been moved as suggested.

How do we prevent ComputeReductionResult from being treated as dead code, if-not/until hooked-up to a LiveOut.

At the moment, ComputeReductionResult needs to be considered as having side effects to prevent DCE.

[ayalz]

Must the reduction chain now terminate at FinalReductionResult? I.e., assert that UserRecipe is non-null by (non-dyn) casting U into a recipe.

[fhahn]

At this point yes, updated to use cast, thanks!

[ayalz]

Would it be better to introduce these Select's for non in-loop reductions under tail folding, along with the introduction of ComputeReductionResult? Rather than first introducing the latter and then searching to replace its operand by the former.

Note that these Selects originally had to be introduced inside the loop, but can now be placed in the middle block alongside ComputeReductionResult.

nit (unrelated to this patch): is the enclosing for statement mis-indented, along with its preceding Builder.setInsertPoint(&*LatchVPBB->begin()); statement?

[ayalz]

Encapsulated approach is preferred?

[ayalz]

Various comments, nice to see generating bc.merge.rdx fits in collectMergePhiForReduction().

[ayalz]

Should Builder support direct insertion of a built recipe before its insertion point?

[fhahn]

Adjusted independently in 6dda74c, will rebase and resolve the conflicts here once approved.

[ayalz]

Above early-continue for in-loop reductions has been lifted, but assert below indicates checking in-loop is redundant for truncated reductions? Otherwise assert is redundant (and probably has been).

[fhahn]

It's now redundant, removed, thanks!

[ayalz]

Should the assert and original check be kept intact, i.e., avoid introducing the check if !PhiR->isInLoop()?
See collectInLoopReductions():

    // We don't collect reductions that are type promoted (yet).
    if (RdxDesc.getRecurrenceType() != Phi->getType())
      continue;

going back to 745bf6c

[fhahn]

Looks like the assert isn't triggered, replaced the InLoop check with assert again.

[ayalz]

Instead of existing RAUW + excluding Trunc to keep use Result/ExitValue?

[fhahn]

Yes, this way preserves other users of OrigExitValue in all cases

[ayalz]

Red is now aka OrigExitValue

[fhahn]

Yep, updated, thanks!

[ayalz]

Could be an RAUW except for the relevant user set somewhere above, which could be recorded and retained.

[fhahn]

Yes, mainly we don't want to update the reduction phi user.

[ayalz]

Suggested change

	// Check if \p RedResult is a CompueReductionResult instruction, and if it is
	// Check if \p RedResult is a ComputeReductionResult instruction, and if it is

[fhahn]

Done, thanks!

[ayalz]

Should the assert and original check be kept intact, i.e., avoid introducing the check if !PhiR->isInLoop()?
See collectInLoopReductions():

    // We don't collect reductions that are type promoted (yet).
    if (RdxDesc.getRecurrenceType() != Phi->getType())
      continue;

going back to 745bf6c

[ayalz]

Current option is ok, temporarily. Alternatives listed have their disadvantages.

[ayalz]

The second, true operand %and.red.next is probably printed as well, just not CHECK'ed.

[ayalz]

Very well!

[ayalz]

Good! removeLiveOut() is still needed for other header-phi live-outs?

[ayalz]

Looks good to me! Adding final nits.

[ayalz]

nit (unrelated to this patch): call Builder.createSelect()?

[fhahn]

Will rebase the patch before landing, which will remove those changes.

[ayalz]

nit: CHECK additional operand here too.

[fhahn]

Will update before landing, thanks!

[ayalz]

nit: CHECK additional operand here too.

[fhahn]

Will update before landing, thanks!

[ayalz]

nit: CHECK additional operand here too.

[fhahn]

Will update before landing, thanks!