[VPlan] Track both incoming values for first-order recurrence phis.

[fhahn]

Summary

This patch updates VPWidenPHI recipes for first-order recurrences to
also track the incoming value from the back-edge. Similar to D99294,
which did the same for reductions.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4760

This does raise an eyebrow... rename StartVPV and StartV to something like ReductionStartVPV and ReductionStartV? Only start values of reductions are handled here, those of firstOrderRecurrences are handled in fixFirstOrderRecurrence().

Better yet to first handle firstOrderRecurrences, separately? It's just creating the phi (which will eventually be replaced by the one fixFirstOrderRecurrence() creates) and setting the state, per part.

8976

This StartV for non-reduction phi's is currently not being used, except perhaps as a placeholder to ensure backedge takes 2nd place, right? Should potentially serve fixFirstOrderRecurrence()?

8977

What's left as 'else' here, pointer inductions? Worth a TODO to model their backedges as well?

8980

dyn_cast may return null; use cast<>?

8982

Comments in VPRecipeBuilder.h that PhisToFix handles reduction phi's only should be updated, to include FOR phi's as well:

/// Cross-iteration **reduction** phis for which we need to add the incoming value
/// from the backedge after all recipes have been created.
SmallVector<VPWidenPHIRecipe *, 4> PhisToFix;

/// Add the incoming values from the backedge to **reduction** cross-iteration
/// phis.
void fixHeaderPhis();

Comments in VPlan.h about the backedge of VPWidenPHIRecipe should be updated as well:

/// For **reduction** PHIs, RdxDesc must point to the corresponding recurrence
/// descriptor, the start value is the first operand of the recipe and the
/// incoming value from the backedge is the second operand.

/// Returns the incoming value from the loop backedge, if it is a **reduction**.
VPValue *getBackedgeValue()

[fhahn]

Addressed comments, thanks!

I also updated fixFirstOrderRecurrence to use the VPValues for the start and backedge values.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4760

I re-organized the code a bit more. WDYT?

8976

I updated the patch to use both the start and backedge VPValue in fixFirstOrderRecurrence.

8977

Added the todo.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4772–4806

(Unrelated to this patch, but setting the loop's upper bound in advance to unsigned LastPartForNewPhi = IsOrdered ? 1 : State.UF; seems clearer than breaking inside the loop. Here and below.)

4823

Instead of asking if (RdxDesc) here, move the following block earlier to after FOR(P)'s return; i.e., have the "if (RdxDesc || FOR(P))" fully handle both?

This will also allow to keep the definitions of ScalarPHI and VecTy inside "if (RdxDesc || FOR(P))".

8973

Add a VPWidenPHIRecipe(Phi, *StartV) constructor? (If not split into VPWidenReductionPHIRecipe, VPWidenFORPHIRecipe, VPWidenPointerIVPHIRecipe.)

[fhahn]

Address latest comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4772–4806

Done in 7f369819774d

4823

Sounds good, I moved the code up.

8973

I added the new constructor. I'm planning to split them up as discussed at the other phi related patches as follow ups.

[fhahn]

Use LastPartForNewPhi in second loop as well.

[Ayal]

This looks good to me, thanks!
Adding a couple of final nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4302

Better place this definition of Phi after the following comment, right before it is needed for defining Start?

4761

Better keep this definition of IsOrdered inside the if (RdxDesc || Legal->isFirstOrderRecurrence(P)) { below, right before it is used to define LastPartForNewPhi?

[fhahn]

Thanks Ayal! I included your suggestions in the committed version.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

ExtractForPhiUsedOutsideLoop = Builder.CreateExtractElement(
         Incoming, Idx, "vector.recur.extract.for.phi");
   } else if (UF > 1)
     // When loop is unrolled without vectorizing, initialize
     // ExtractForPhiUsedOutsideLoop with the value just prior to unrolled value
     // of `Incoming`. This is analogous to the vectorized case above: extracting
     // the second last element when VF > 1.
     ExtractForPhiUsedOutsideLoop = State.get(PreviousDef, UF - 2);
 
   // Fix the initial value of the original recurrence in the scalar loop.
   Builder.SetInsertPoint(&*LoopScalarPreHeader->begin());
   auto *Start = Builder.CreatePHI(Phi->getType(), 2, "scalar.recur.init");
   for (auto *BB : predecessors(LoopScalarPreHeader)) {
     auto *Incoming = BB == LoopMiddleBlock ? ExtractForScalar : ScalarInit;
     Start->addIncoming(Incoming, BB);
   }
 
   Phi->setIncomingValueForBlock(LoopScalarPreHeader, Start);
if (EnableVPlanNativePath) {
     OrigPHIsToFix.push_back(P);
 
     return;
   }
 
   assert(PN->getParent() == OrigLoop->getHeader() &&
          "Non-header phis should have been handled elsewhere");
 
-  VPValue *StartVPV = PhiR->getStartValue();
+  VPValue *StartVPV = RdxDesc ? PhiR->getStartValue() : nullptr;
   Value *StartV = StartVPV ? StartVPV->getLiveInIRValue() : nullptr;
   // In order to support recurrences we need to be able to vectorize Phi nodes.
   // Phi nodes have cycles, so we need to vectorize them in two stages. This is
   // stage #1: We create a new vector PHI node with no incoming edges. We'll use
   // this value when we vectorize all of the instructions that use the PHI.
   if (RdxDesc || Legal->isFirstOrderRecurrence(P)) {
     Value *Iden = nullptr;
     bool ScalarPHI =
         (State.VF.isScalar()) || Cost->isInLoopReduction(cast<PHINode>(PN));
     Type *VecTy =
         ScalarPHI ? PN->getType() : VectorType::get(PN->getType(), State.VF);
 
     if (RdxDesc) {
       assert(Legal->isReductionVariable(P) && StartV &&
              "RdxDesc should only be set for reduction variables; in that case "
              "a StartV is also required");
       RecurKind RK = RdxDesc->getRecurrenceKind();
       if (RecurrenceDescriptor::isMinMaxRecurrenceKind(RK)) {
         // MinMax reduction have the start value as their identify.
         if (ScalarPHI) {
           Iden = StartV;
         } else {
           IRBuilderBase::InsertPointGuard IPBuilder(Builder);
           Builder.SetInsertPoint(LoopVectorPreHeader->getTerminator());
           StartV = Iden =
               Builder.CreateVectorSplat(State.VF, StartV, "minmax.ident");
         }
       } else {
         Constant *IdenC = RecurrenceDescriptor::getRecurrenceIdentity(
             RK, VecTy->getScalarType(), RdxDesc->getFastMathFlags());
         Iden = IdenC;
 
         if (!ScalarPHI) {
           Iden = ConstantVector::getSplat(State.VF, IdenC);
           IRBuilderBase::InsertPointGuard IPBuilder(Builder);
           Builder.SetInsertPoint(LoopVectorPreHeader->getTerminator());
           Constant *Zero = Builder.getInt32(0);
           StartV = Builder.CreateInsertElement(Iden, StartV, Zero);
         }
       }
     }
 
     bool IsOrdered = State.VF.isVector() &&
                      Cost->isInLoopReduction(cast<PHINode>(PN)) &&
                      Cost->useOrderedReductions(*RdxDesc);
 
     for (unsigned Part = 0; Part < State.UF; ++Part) {
       // This is phase one of vectorizing PHIs.
       if (Part > 0 && IsOrdered)
         return;
       Value *EntryPart = PHINode::Create(
           VecTy, 2, "vec.phi", &*LoopVectorBody->getFirstInsertionPt());
       State.set(PhiR, EntryPart, Part);
       if (StartV) {
         // Make sure to add the reduction start value only to the
         // first unroll part.
         Value *StartVal = (Part == 0) ? StartV : Iden;
         cast<PHINode>(EntryPart)->addIncoming(StartVal, LoopVectorPreHeader);
       }
     }
     return;
   }
 
   assert(!Legal->isReductionVariable(P) &&
          "reductions should be handled above");
 
   setDebugLocFromInst(Builder, P);
 
   // This PHINode must be an induction variable.
   // Make sure that we know about it.
   assert(Legal->getInductionVars().count(P) && "Not an induction variable");
VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
 
   VPRecipeBase *Recipe;
   if (auto Phi = dyn_cast<PHINode>(Instr)) {
     if (Phi->getParent() != OrigLoop->getHeader())
       return tryToBlend(Phi, Operands, Plan);
     if ((Recipe = tryToOptimizeInductionPHI(Phi, Operands)))
       return toVPRecipeResult(Recipe);
 
+    VPWidenPHIRecipe *PhiRecipe = nullptr;
+    VPValue *StartV = Operands[0];
     if (Legal->isReductionVariable(Phi)) {
       RecurrenceDescriptor &RdxDesc = Legal->getReductionVars()[Phi];
       assert(RdxDesc.getRecurrenceStartValue() ==
              Phi->getIncomingValueForBlock(OrigLoop->getLoopPreheader()));
-      VPValue *StartV = Operands[0];
 
-      auto *PhiRecipe = new VPWidenPHIRecipe(Phi, RdxDesc, *StartV);
-      PhisToFix.push_back(PhiRecipe);
+      PhiRecipe = new VPWidenPHIRecipe(Phi, RdxDesc, *StartV);
       // Record the incoming value from the backedge, so we can add the incoming
       // value from the backedge after all recipes have been created.
       recordRecipeOf(cast<Instruction>(
           Phi->getIncomingValueForBlock(OrigLoop->getLoopLatch())));
-      return toVPRecipeResult(PhiRecipe);
+      PhisToFix.push_back(PhiRecipe);
+    } else {
+      PhiRecipe = new VPWidenPHIRecipe(Phi);
+      PhiRecipe->addOperand(StartV);
+      if (Legal->isFirstOrderRecurrence(Phi)) {
+        Value *Incoming =
+            Phi->getIncomingValueForBlock(OrigLoop->getLoopLatch());
+        Instruction *IncI = dyn_cast<Instruction>(Incoming);
+        recordRecipeOf(IncI);
+        PhisToFix.push_back(PhiRecipe);
+      }
     }
 
-    return toVPRecipeResult(new VPWidenPHIRecipe(Phi));
+    return toVPRecipeResult(PhiRecipe);
   }
 
   if (isa<TruncInst>(Instr) &&
       (Recipe = tryToOptimizeInductionTruncate(cast<TruncInst>(Instr), Operands,
                                                Range, *Plan)))
     return toVPRecipeResult(Recipe);
 
   if (!shouldWiden(Instr, Range))

llvm/test/Transforms/LoopVectorize/first-order-recurrence-sink-replicate-region.ll

 ; REQUIRES: asserts
 ; RUN: opt < %s -loop-vectorize -force-vector-width=2 -force-vector-interleave=1 -debug-only=loop-vectorize 2>&1 | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 
 ; Test cases for PR50009, which require sinking a replicate-region due to a
 ; first-order recurrence.
 
 define void @sink_replicate_region_1(i32 %x, i8* %ptr) optsize {
 ; CHECK-LABEL: sink_replicate_region_1
 ; CHECK:      VPlan 'Initial VPlan for VF={2},UF>=1' {
 ; CHECK-NEXT: loop:
-; CHECK-NEXT:   WIDEN-PHI %0 = phi 0, %conv
+; CHECK-NEXT:   WIDEN-PHI ir<%0> = phi ir<0>, ir<%conv>
 ; CHECK-NEXT:   WIDEN-INDUCTION %iv = phi 0, %iv.next
 ; CHECK-NEXT:   EMIT vp<%3> = icmp ule ir<%iv> vp<%0>
 ; CHECK-NEXT: Successor(s): loop.0
 ; CHECK-EMPTY:
 ; CHECK-NEXT: loop.0:
 ; CHECK-NEXT: Successor(s): pred.load
 ; CHECK-EMPTY:
 ; CHECK-NEXT: <xVFxUF> pred.load: {

 exit:
   ret void
 }
 
 define void @sink_replicate_region_2(i32 %x, i8 %y, i32* %ptr) optsize {
 ; CHECK-LABEL: sink_replicate_region_2
 ; CHECK:      VPlan 'Initial VPlan for VF={2},UF>=1' {
 ; CHECK-NEXT: loop:
-; CHECK-NEXT:   WIDEN-PHI %recur = phi 0, %recur.next
+; CHECK-NEXT:   WIDEN-PHI ir<%recur> = phi ir<0>, ir<%recur.next>
 ; CHECK-NEXT:   WIDEN-INDUCTION %iv = phi 0, %iv.next
 ; CHECK-NEXT:   EMIT vp<%3> = icmp ule ir<%iv> vp<%0>
 ; CHECK-NEXT: Successor(s): loop.0
 ; CHECK-EMPTY:
 ; CHECK-NEXT: loop.0:
 ; CHECK-NEXT:   WIDEN ir<%recur.next> = sext ir<%y>
 ; CHECK-NEXT: Successor(s): pred.srem
 ; CHECK-EMPTY:

 exit:
   ret void
 }
 
 define i32 @sink_replicate_region_3_reduction(i32 %x, i8 %y, i32* %ptr) optsize {
 ; CHECK-LABEL: sink_replicate_region_3_reduction
 ; CHECK:      VPlan 'Initial VPlan for VF={2},UF>=1' {
 ; CHECK-NEXT: loop:
-; CHECK-NEXT:   WIDEN-PHI %recur = phi 0, %recur.next
+; CHECK-NEXT:   WIDEN-PHI ir<%recur> = phi ir<0>, ir<%recur.next>
 ; CHECK-NEXT:   WIDEN-INDUCTION %iv = phi 0, %iv.next
 ; CHECK-NEXT:   WIDEN-PHI ir<%and.red> = phi ir<1234>, ir<%and.red.next>
 ; CHECK-NEXT:   EMIT vp<%4> = icmp ule ir<%iv> vp<%0>
 ; CHECK-NEXT: Successor(s): loop.0
 ; CHECK-EMPTY:
 ; CHECK-NEXT: loop.0:
 ; CHECK-NEXT:   WIDEN ir<%recur.next> = sext ir<%y>
 ; CHECK-NEXT: Successor(s): pred.srem

 }
 
 ; To sink the replicate region containing %rem, we need to split the block
 ; containing %conv at the end, because %conv is the last recipe in the block.
 define void @sink_replicate_region_4_requires_split_at_end_of_block(i32 %x, i8* %ptr) optsize {
 ; CHECK-LABEL: sink_replicate_region_4_requires_split_at_end_of_block
 ; CHECK:      VPlan 'Initial VPlan for VF={2},UF>=1' {
 ; CHECK-NEXT: loop:
-; CHECK-NEXT:   WIDEN-PHI %0 = phi 0, %conv
+; CHECK-NEXT:   WIDEN-PHI ir<%0> = phi ir<0>, ir<%conv>
 ; CHECK-NEXT:   WIDEN-INDUCTION %iv = phi 0, %iv.next
 ; CHECK-NEXT:   EMIT vp<%3> = icmp ule ir<%iv> vp<%0>
 ; CHECK-NEXT:   REPLICATE ir<%gep> = getelementptr ir<%ptr>, ir<%iv>
 ; CHECK-NEXT: Successor(s): loop.0
 ; CHECK-EMPTY:
 ; CHECK-NEXT: loop.0:
 ; CHECK-NEXT: Successor(s): pred.load
 ; CHECK-EMPTY: