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

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2023-04-10T13:02:44+01:00
New Revision: c255eb2c4bcc72b2f3fb3ed8f38df2b6f6978d2a

URL: https://github.com/llvm/llvm-project/commit/c255eb2c4bcc72b2f3fb3ed8f38df2b6f6978d2a
DIFF: https://github.com/llvm/llvm-project/commit/c255eb2c4bcc72b2f3fb3ed8f38df2b6f6978d2a.diff

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

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.

Reviewed By: Ayal, michaelmaitland

Differential Revision: https://reviews.llvm.org/D147472

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
    llvm/lib/Transforms/Vectorize/VPlan.h

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 2f97a02cc9e1d..3a88535bee911 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -3900,24 +3900,22 @@ void InnerLoopVectorizer::fixFixedOrderRecurrence(
   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)) {
-    SmallPtrSet<PHINode *, 2> ToFix;
-    for (User *U : Phi->users())
-      if (isa<PHINode>(U) && cast<Instruction>(U)->getParent() == LoopExitBlock)
-        ToFix.insert(cast<PHINode>(U));
-    for (PHINode *LCSSAPhi : ToFix) {
-      LCSSAPhi->addIncoming(ExtractForPhiUsedOutsideLoop, LoopMiddleBlock);
-      State.Plan->removeLiveOut(LCSSAPhi);
-    }
+  auto RecurSplice = cast<VPInstruction>(*PhiR->user_begin());
+  assert(PhiR->getNumUsers() == 1 &&
+         RecurSplice->getOpcode() ==
+             VPInstruction::FirstOrderRecurrenceSplice &&
+         "recurrence phi must have a single user: FirstOrderRecurrenceSplice");
+  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) &&
+           "no live-outs expected when scalar epilogue is required");
+    PHINode *LCSSAPhi = LiveOut->getPhi();
+    LCSSAPhi->addIncoming(ExtractForPhiUsedOutsideLoop, LoopMiddleBlock);
+    State.Plan->removeLiveOut(LCSSAPhi);
   }
 }
 

diff  --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index b84022c14cce1..1dcc5d194eba2 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -666,6 +666,10 @@ class VPLiveOut : public VPUser {
   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