!fixup address latest comments, thansk!

fhahn · fhahn · commit f028ab3f432d · 2024-07-09T14:45:23.000+01:00
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -3290,17 +3290,6 @@ void InnerLoopVectorizer::fixVectorizedLoop(VPTransformState &State,
   if (EnableVPlanNativePath)
     fixNonInductionPHIs(Plan, State);
 
-  // At this point every instruction in the original loop is widened to a
-  // vector form. Note that fixing reduction phis, as well as extracting the
-  // exit and resume values for fixed-order recurrences are already modeled in
-  // VPlan. All that remains to do here is to create a phi in the scalar
-  // pre-header for each fixed-order recurrence resume value.
-  // TODO: Also model creating phis in the scalar pre-header in VPlan.
-  for (const auto &[_, LO] : to_vector(Plan.getLiveOuts())) {
-    if (!Legal->isFixedOrderRecurrence(LO->getPhi()))
-      continue;
-  }
-
   // Forget the original basic block.
   PSE.getSE()->forgetLoop(OrigLoop);
   PSE.getSE()->forgetBlockAndLoopDispositions();
@@ -3337,8 +3326,6 @@ void InnerLoopVectorizer::fixVectorizedLoop(VPTransformState &State,
                    VectorLoop->getHeader(), Plan, State);
   }
 
-  // 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,
                                State.CFG.ExitBB->getFirstNonPHIIt());
   for (const auto &KV : Plan.getLiveOuts())
@@ -8485,9 +8472,59 @@ static void addUsersInExitBlock(VPBasicBlock *HeaderVPBB, Loop *OrigLoop,
     Value *IncomingValue =
         ExitPhi.getIncomingValueForBlock(ExitingBB);
     VPValue *V = Builder.getVPValueOrAddLiveIn(IncomingValue, Plan);
-    Plan.addLiveOut(
-        &ExitPhi, V,
-        cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getSingleSuccessor()));
+    Plan.addLiveOut(&ExitPhi, V);
+  }
+}
+
+/// Feed a resume value for every FOR from the vector loop to the scalar loop,
+/// if middle block branches to scalar preheader, by introducing ExtractFromEnd
+/// and ResumePhi recipes in each, respectively, and a VPLiveOut which uses the
+/// latter and corresponds to the scalar header.
+static void addLiveOutsForFirstOrderRecurrences(VPlan &Plan) {
+  VPRegionBlock *VectorRegion = Plan.getVectorLoopRegion();
+
+  // Start by finding out if middle block branches to scalar preheader.
+  // TODO: Should be replaced by
+  // Plan->getScalarLoopRegion()->getSinglePredecessor() in the future once the
+  // scalar region is modeled as well.
+  VPBasicBlock *ScalarPHVPBB = nullptr;
+  auto *MiddleVPBB = cast<VPBasicBlock>(VectorRegion->getSingleSuccessor());
+  for (VPBlockBase *Succ : MiddleVPBB->getSuccessors()) {
+    if (isa<VPIRBasicBlock>(Succ))
+      continue;
+    assert(!ScalarPHVPBB && "Two candidates for ScalarPHVPBB?");
+    ScalarPHVPBB = cast<VPBasicBlock>(Succ);
+  }
+  if (!ScalarPHVPBB)
+    return;
+
+  for (auto &H : VectorRegion->getEntryBasicBlock()->phis()) {
+    auto *FOR = dyn_cast<VPFirstOrderRecurrencePHIRecipe>(&H);
+    if (!FOR)
+      continue;
+
+    VPBuilder B(ScalarPHVPBB);
+    VPBuilder MiddleBuilder(MiddleVPBB);
+    // Reset insert point so new recipes are inserted before terminator and
+    // condition, if there is either the former or both.
+    if (auto *Terminator = MiddleVPBB->getTerminator()) {
+      auto *Condition = dyn_cast<VPInstruction>(Terminator->getOperand(0));
+      assert((!Condition || Condition->getParent() == MiddleVPBB) &&
+             "Condition expected in MiddleVPBB");
+      MiddleBuilder.setInsertPoint(Condition ? Condition : Terminator);
+    }
+
+    // Extract the resume value and create a new VPLiveOut for it.
+    auto *Resume = MiddleBuilder.createNaryOp(
+        VPInstruction::ExtractFromEnd,
+        {FOR->getBackedgeValue(),
+         Plan.getOrAddLiveIn(
+             ConstantInt::get(Plan.getCanonicalIV()->getScalarType(), 1))},
+        {}, "vector.recur.extract");
+    auto *ResumePhiRecipe =
+        B.createNaryOp(VPInstruction::ResumePhi, {Resume, FOR->getStartValue()},
+                       {}, "scalar.recur.init");
+    Plan.addLiveOut(cast<PHINode>(FOR->getUnderlyingInstr()), ResumePhiRecipe);
   }
 }
 
@@ -8655,48 +8692,7 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
          "VPBasicBlock");
   RecipeBuilder.fixHeaderPhis();
 
-  auto *MiddleVPBB =
-      cast<VPBasicBlock>(Plan->getVectorLoopRegion()->getSingleSuccessor());
-
-  VPBasicBlock *ScalarPH = nullptr;
-  for (VPBlockBase *Succ : MiddleVPBB->getSuccessors()) {
-    auto *VPBB = dyn_cast<VPBasicBlock>(Succ);
-    if (VPBB && !isa<VPIRBasicBlock>(VPBB)) {
-      ScalarPH = VPBB;
-      break;
-    }
-  }
-
-  if (ScalarPH) {
-    for (auto &H : HeaderVPBB->phis()) {
-      auto *FOR = dyn_cast<VPFirstOrderRecurrencePHIRecipe>(&H);
-      if (!FOR)
-        continue;
-      VPBuilder B(ScalarPH);
-      VPBuilder MiddleBuilder;
-      // Set insert point so new recipes are inserted before terminator and
-      // condition, if there is either the former or both.
-      if (MiddleVPBB->getNumSuccessors() != 2)
-        MiddleBuilder.setInsertPoint(MiddleVPBB);
-      else if (isa<VPInstruction>(MiddleVPBB->getTerminator()->getOperand(0)))
-        MiddleBuilder.setInsertPoint(
-            &*std::prev(MiddleVPBB->getTerminator()->getIterator()));
-      else
-        MiddleBuilder.setInsertPoint(MiddleVPBB->getTerminator());
-
-      // Extract the resume value and create a new VPLiveOut for it.
-      auto *Resume = MiddleBuilder.createNaryOp(
-          VPInstruction::ExtractFromEnd,
-          {FOR->getBackedgeValue(),
-           Plan->getOrAddLiveIn(
-               ConstantInt::get(Plan->getCanonicalIV()->getScalarType(), 1))},
-          {}, "vector.recur.extract");
-      auto *R =
-          B.createNaryOp(VPInstruction::ExitPhi, {Resume, FOR->getStartValue()},
-                         {}, "scalar.recur.init");
-      Plan->addLiveOut(cast<PHINode>(FOR->getUnderlyingInstr()), R, ScalarPH);
-    }
-  }
+  addLiveOutsForFirstOrderRecurrences(*Plan);
 
   // ---------------------------------------------------------------------------
   // Transform initial VPlan: Apply previously taken decisions, in order, to
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -915,11 +915,6 @@ void VPlan::prepareToExecute(Value *TripCountV, Value *VectorTripCountV,
 /// VPBB, if any, are rewired to the new VPIRBasicBlock.
 static void replaceVPBBWithIRVPBB(VPBasicBlock *VPBB, BasicBlock *IRBB) {
   VPIRBasicBlock *IRMiddleVPBB = new VPIRBasicBlock(IRBB);
-  for (auto &[_, LO] : VPBB->getPlan()->getLiveOuts()) {
-    if (LO->getPred() == VPBB)
-      LO->setPred(IRMiddleVPBB);
-  }
-
   for (auto &R : make_early_inc_range(*VPBB))
     R.moveBefore(*IRMiddleVPBB, IRMiddleVPBB->end());
   VPBlockBase *PredVPBB = VPBB->getSinglePredecessor();
@@ -1135,9 +1130,9 @@ LLVM_DUMP_METHOD
 void VPlan::dump() const { print(dbgs()); }
 #endif
 
-void VPlan::addLiveOut(PHINode *PN, VPValue *V, VPBasicBlock *Pred) {
+void VPlan::addLiveOut(PHINode *PN, VPValue *V) {
   assert(LiveOuts.count(PN) == 0 && "an exit value for PN already exists");
-  LiveOuts.insert({PN, new VPLiveOut(PN, V, Pred)});
+  LiveOuts.insert({PN, new VPLiveOut(PN, V)});
 }
 
 static void remapOperands(VPBlockBase *Entry, VPBlockBase *NewEntry,
@@ -1206,18 +1201,9 @@ VPlan *VPlan::duplicate() {
   remapOperands(Preheader, NewPreheader, Old2NewVPValues);
   remapOperands(Entry, NewEntry, Old2NewVPValues);
 
-  DenseMap<VPBlockBase *, VPBlockBase *> Old2NewVPBlocks;
-  VPBlockBase *OldMiddle = getVectorLoopRegion()->getSingleSuccessor();
-  VPBlockBase *NewMiddle = NewPlan->getVectorLoopRegion()->getSingleSuccessor();
-  Old2NewVPBlocks[OldMiddle] = NewMiddle;
-  for (const auto &[Old, New] :
-       zip(OldMiddle->getSuccessors(), NewMiddle->getSuccessors()))
-    Old2NewVPBlocks[Old] = New;
-
   // Clone live-outs.
   for (const auto &[_, LO] : LiveOuts)
-    NewPlan->addLiveOut(LO->getPhi(), Old2NewVPValues[LO->getOperand(0)],
-                        cast<VPBasicBlock>(Old2NewVPBlocks[LO->getPred()]));
+    NewPlan->addLiveOut(LO->getPhi(), Old2NewVPValues[LO->getOperand(0)]);
 
   // Initialize remaining fields of cloned VPlan.
   NewPlan->VFs = VFs;
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -682,27 +682,23 @@ class VPBlockBase {
 };
 
 /// A value that is used outside the VPlan. The operand of the user needs to be
-/// added to the associated LCSSA phi node.
+/// added to the associated LCSSA phi node. The incoming block from VPlan is
+/// determined by where the VPValue is defined: if it is defined by a recipe
+/// outside a region, its parent block is used, otherwise the middle block is
+/// used.
 class VPLiveOut : public VPUser {
   PHINode *Phi;
 
-  /// Predecessor in VPlan of this live-out. Used to as block to set the
-  /// incoming value for.
-  VPBasicBlock *Pred;
-
 public:
-  VPLiveOut(PHINode *Phi, VPValue *Op, VPBasicBlock *Pred)
-      : VPUser({Op}, VPUser::VPUserID::LiveOut), Phi(Phi), Pred(Pred) {}
+  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.
+  /// Fixup the wrapped LCSSA phi node. This means we need to add the
+  /// appropriate incoming value from the precessor Pred.
   void fixPhi(VPlan &Plan, VPTransformState &State);
 
   /// Returns true if the VPLiveOut uses scalars of operand \p Op.
@@ -714,11 +710,6 @@ class VPLiveOut : public VPUser {
 
   PHINode *getPhi() const { return Phi; }
 
-  /// Returns to incoming block for which to set the value.
-  VPBasicBlock *getPred() const { return Pred; }
-
-  void setPred(VPBasicBlock *Pred) { this->Pred = Pred; }
-
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the VPLiveOut to \p O.
   void print(raw_ostream &O, VPSlotTracker &SlotTracker) const;
@@ -1202,8 +1193,9 @@ class VPInstruction : public VPRecipeWithIRFlags {
     ExplicitVectorLength,
     /// Creates a scalar phi in a leaf VPBB with a single predecessor in VPlan.
     /// The first operand is the incoming value from the predecessor in VPlan,
-    /// the second operand is the incoming value for all other predecessors.
-    ExitPhi,
+    /// the second operand is the incoming value for all other predecessors
+    /// (which are currently not modeled in VPlan).
+    ResumePhi,
     CalculateTripCountMinusVF,
     // Increment the canonical IV separately for each unrolled part.
     CanonicalIVIncrementForPart,
@@ -3346,7 +3338,7 @@ class VPlan {
     return cast<VPCanonicalIVPHIRecipe>(&*EntryVPBB->begin());
   }
 
-  void addLiveOut(PHINode *PN, VPValue *V, VPBasicBlock *Pred);
+  void addLiveOut(PHINode *PN, VPValue *V);
 
   void removeLiveOut(PHINode *PN) {
     delete LiveOuts[PN];
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -193,7 +193,13 @@ void VPLiveOut::fixPhi(VPlan &Plan, VPTransformState &State) {
   auto Lane = vputils::isUniformAfterVectorization(ExitValue)
                   ? VPLane::getFirstLane()
                   : VPLane::getLastLaneForVF(State.VF);
-  BasicBlock *PredBB = State.CFG.VPBB2IRBB[Pred];
+  VPBasicBlock *PredVPBB =
+      cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getSingleSuccessor());
+  VPRecipeBase *DefRecipe = ExitValue->getDefiningRecipe();
+  if (DefRecipe && !DefRecipe->getParent()->getParent())
+    PredVPBB = DefRecipe->getParent();
+  BasicBlock *PredBB = State.CFG.VPBB2IRBB[PredVPBB];
+  State.Builder.SetInsertPoint(PredBB, PredBB->getFirstNonPHIIt());
   Value *V = State.get(ExitValue, VPIteration(State.UF - 1, Lane));
   if (Phi->getBasicBlockIndex(PredBB) != -1)
     Phi->setIncomingValueForBlock(PredBB, V);
@@ -304,7 +310,7 @@ bool VPInstruction::canGenerateScalarForFirstLane() const {
   case VPInstruction::CanonicalIVIncrementForPart:
   case VPInstruction::PtrAdd:
   case VPInstruction::ExplicitVectorLength:
-  case VPInstruction::ExitPhi:
+  case VPInstruction::ResumePhi:
     return true;
   default:
     return false;
@@ -595,20 +601,24 @@ Value *VPInstruction::generatePerPart(VPTransformState &State, unsigned Part) {
     Value *Addend = State.get(getOperand(1), Part, /* IsScalar */ true);
     return Builder.CreatePtrAdd(Ptr, Addend, Name);
   }
-  case VPInstruction::ExitPhi: {
+  case VPInstruction::ResumePhi: {
     if (Part != 0)
       return State.get(this, 0, /*IsScalar*/ true);
     Value *IncomingFromVPlanPred =
         State.get(getOperand(0), Part, /* IsScalar */ true);
-    Value *IncomingFromOtherPred =
+    Value *IncomingFromOtherPreds =
         State.get(getOperand(1), Part, /* IsScalar */ true);
-    auto *NewPhi = Builder.CreatePHI(IncomingFromOtherPred->getType(), 2, Name);
+    auto *NewPhi =
+        Builder.CreatePHI(IncomingFromOtherPreds->getType(), 2, Name);
     BasicBlock *VPlanPred =
         State.CFG
             .VPBB2IRBB[cast<VPBasicBlock>(getParent()->getSinglePredecessor())];
     NewPhi->addIncoming(IncomingFromVPlanPred, VPlanPred);
-    for (auto *BB : predecessors(Builder.GetInsertBlock()))
-      NewPhi->addIncoming(IncomingFromOtherPred, BB);
+    for (auto *OtherPred : predecessors(Builder.GetInsertBlock())) {
+      assert(OtherPred != VPlanPred &&
+             "VPlan predecessors should not be connected yet");
+      NewPhi->addIncoming(IncomingFromOtherPreds, OtherPred);
+    }
     return NewPhi;
   }
 
@@ -619,7 +629,7 @@ Value *VPInstruction::generatePerPart(VPTransformState &State, unsigned Part) {
 
 bool VPInstruction::isVectorToScalar() const {
   return getOpcode() == VPInstruction::ExtractFromEnd ||
-         getOpcode() == VPInstruction::ExitPhi ||
+         getOpcode() == VPInstruction::ResumePhi ||
          getOpcode() == VPInstruction::ComputeReductionResult;
 }
 
@@ -749,7 +759,7 @@ void VPInstruction::print(raw_ostream &O, const Twine &Indent,
   case VPInstruction::ActiveLaneMask:
     O << "active lane mask";
     break;
-  case VPInstruction::ExitPhi:
+  case VPInstruction::ResumePhi:
     O << "exit-phi";
     break;
   case VPInstruction::ExplicitVectorLength:
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -937,7 +937,9 @@ bool VPlanTransforms::adjustFixedOrderRecurrences(VPlan &Plan,
     Type *IntTy = Plan.getCanonicalIV()->getScalarType();
 
     // Extract the penultimate value of the recurrence and update VPLiveOut
-    // users of the recurrence splice.
+    // users of the recurrence splice. Note that the extract of the final value
+    // used to resume in the scalar loop is created earlier during VPlan
+    // construction.
     auto *Penultimate = cast<VPInstruction>(MiddleBuilder.createNaryOp(
         VPInstruction::ExtractFromEnd,
         {FOR->getBackedgeValue(),
