[VPlan] Use VPLiveOut to update FOR live-out users.

[fhahn]

Summary

Instead of iterating over all LCSSA phis in the exit block, collect all
LiveOut users of the FOR splice VPInstruction and only update those
users.

Building on top of D147471, this removes an access to the cost model
after VPlan execution.

Depends on D147471.

[michaelmaitland]

LGTM.

[fhahn]

Rebase after recent upstream changes.

[Ayal]

Looks good to me, and raises some thoughts ;-)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3859

nit (Independent of this patch): else block deserves parentheses considering comments?

3879

nit: "only user should be FirstOrderRecurrenceSplice VPInstruction" >>
"Recurrence Phi must have a single user: FirstOrderRecurrenceSplice"?

3880

nit: worth asserting opcode is FirstOrderRecurrenceSplice ?

3884

(Independent of this patch:) Do multiple live-out users feed multiple LCSSAPhi's, should they all feed the same LCSSAPhi, are there tests covering multiple live-out users?

3887

nit: add message to assert.

3890

TODO (Independent of this patch): LiveOuts is presumably needed only for delayed destruction of VPLiveOuts, as adding an incoming IR Value to LCSSAPhi should not interfere with traversing the VPUsers of RecurSplice. Why are VPLiveOut VPUsers destroyed here, rather than during the destruction of VPlan, along with all other VPUsers?
This applies to the other cases of removeLiveOut() as well.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3879

adjusted thanks!

3880

Added the check to the assert.

3884

They feed different LCSSA phis, the test for that is llvm/test/Transforms/LoopVectorize/pr36983.ll.

3887

Thanks, I added a message to the assert.

3890

At the moment, logic later relies on any live-outs that cannot fix themselves to be removed so they are excluded from the handling below

  // Fix LCSSA phis not already fixed earlier. Extracts may need to be generated
  // in the exit block, so update the builder.
  State.Builder.SetInsertPoint(State.CFG.ExitBB->getFirstNonPHI());
  for (const auto &KV : Plan.getLiveOuts())
    KV.second->fixPhi(Plan, State);
```.

Once live-outs of FOR & IVs can be handled completely in VPlan, no removal should be necessary.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3890

Thanks for the reminder. LiveOuts that haven't fixed themselves already could alternatively be filtered by checking if their underlying LCSSAPhi has an incoming Value attached; the savings in traversing only those who do comes at the cost of copying ToFix and LiveOuts for delayed destructions, and starting the tear down parts of VPlan early.

This raises some follow-up thoughts - could fixPhi() also take care of fixing LCSSAPhi's of FORs where the LiveOut's operand is a FOR splice, instead of taking care of them here? (Ideally all LiveOut phi's should be fixed at the end of VPlan::execute(), but handling reductions is more involved). There may be an issue of reusing same extract for parallel LCSSAPhi's instead of replicating the same extract for each. This is also relevant for fixing other, non-FOR LCSSAPhis, and relates to the general issue of representing extracts in VPlan. It may be simpler to fuse such parallel LCSSAPhi's together, but those currently reside outside the scope of VPlan.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::fixFixedOrderRecurrence(
   // will be the value when jumping to the exit block from the LoopMiddleBlock,
   // when the scalar loop is not run at all.
   Value *ExtractForPhiUsedOutsideLoop = nullptr;
   if (VF.isVector()) {
     auto *RuntimeVF = getRuntimeVF(Builder, IdxTy, VF);
     auto *Idx = Builder.CreateSub(RuntimeVF, ConstantInt::get(IdxTy, 2));
     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());
   PHINode *Phi = cast<PHINode>(PhiR->getUnderlyingValue());
   auto *Start = Builder.CreatePHI(Phi->getType(), 2, "scalar.recur.init");
   auto *ScalarInit = PhiR->getStartValue()->getLiveInIRValue();
   for (auto *BB : predecessors(LoopScalarPreHeader)) {
     auto *Incoming = BB == LoopMiddleBlock ? ExtractForScalar : ScalarInit;
     Start->addIncoming(Incoming, BB);
   }
 
   Phi->setIncomingValueForBlock(LoopScalarPreHeader, Start);
   Phi->setName("scalar.recur");
 
-  // Finally, fix users of the recurrence outside the loop. The users will need
-  // either the last value of the scalar recurrence or the last value of the
-  // vector recurrence we extracted in the middle block. Since the loop is in
-  // LCSSA form, we just need to find all the phi nodes for the original scalar
-  // recurrence in the exit block, and then add an edge for the middle block.
-  // Note that LCSSA does not imply single entry when the original scalar loop
-  // had multiple exiting edges (as we always run the last iteration in the
-  // scalar epilogue); in that case, there is no edge from middle to exit and
-  // and thus no phis which needed updated.
-  if (!Cost->requiresScalarEpilogue(VF))
-    for (PHINode &LCSSAPhi : LoopExitBlock->phis())
-      if (llvm::is_contained(LCSSAPhi.incoming_values(), Phi)) {
-        LCSSAPhi.addIncoming(ExtractForPhiUsedOutsideLoop, LoopMiddleBlock);
-        State.Plan->removeLiveOut(&LCSSAPhi);
-      }
+  assert(PhiR->getNumUsers() == 1 &&
+         "only user should be FirstOrderRecurrenceSplice VPInstruction");
+  auto RecurSplice = cast<VPInstruction>(*PhiR->user_begin());
+  SmallVector<VPLiveOut *> LiveOuts;
+  for (VPUser *U : RecurSplice->users())
+    if (auto *LiveOut = dyn_cast<VPLiveOut>(U))
+      LiveOuts.push_back(LiveOut);
+
+  for (VPLiveOut *LiveOut : LiveOuts) {
+    assert(!Cost->requiresScalarEpilogue(VF));
+    PHINode *LCSSAPhi = LiveOut->getPhi();
+    LCSSAPhi->addIncoming(ExtractForPhiUsedOutsideLoop, LoopMiddleBlock);
+    State.Plan->removeLiveOut(LCSSAPhi);
 }
 
 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");

llvm/lib/Transforms/Vectorize/VPlan.h

 /// added to the associated LCSSA phi node.
 class VPLiveOut : public VPUser {
   PHINode *Phi;
 
 public:
   VPLiveOut(PHINode *Phi, VPValue *Op)
       : VPUser({Op}, VPUser::VPUserID::LiveOut), Phi(Phi) {}
 
+  static inline bool classof(const VPUser *U) {
+    return U->getVPUserID() == VPUser::VPUserID::LiveOut;
+  }
+
   /// Fixup the wrapped LCSSA phi node in the unique exit block.  This simply
   /// means we need to add the appropriate incoming value from the middle
   /// block as exiting edges from the scalar epilogue loop (if present) are
   /// already in place, and we exit the vector loop exclusively to the middle
   /// block.
   void fixPhi(VPlan &Plan, VPTransformState &State);
 
   /// Returns true if the VPLiveOut uses scalars of operand \p Op.