!fixup address comments, thanks!

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -56,6 +56,7 @@
 #include "llvm/Transforms/Vectorize/LoopVectorize.h"
 #include "LoopVectorizationPlanner.h"
 #include "VPRecipeBuilder.h"
+#include "VPlanPatternMatch.h"
 #include "VPlan.h"
 #include "VPlanAnalysis.h"
 #include "VPlanHCFGBuilder.h"
@@ -606,8 +607,8 @@ class InnerLoopVectorizer {
                     BasicBlock *MiddleBlock, BasicBlock *VectorHeader,
                     VPlan &Plan, VPTransformState &State);
 
-  /// Create the exit value of first order recurrences in the middle block and
-  /// update their users.
+  /// Create the phi node for the resume value of first order recurrences in the scalar preheader and
+  /// update the users in the scalar loop.
   void fixFixedOrderRecurrence(VPLiveOut *LO, VPTransformState &State);
 
   /// Iteratively sink the scalarized operands of a predicated instruction into
@@ -3392,8 +3393,8 @@ void InnerLoopVectorizer::fixVectorizedLoop(VPTransformState &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 creating a phi in the scalar
-  // pre-header for each fixed-rder recurrence resume value.
+  // 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()))
@@ -3473,21 +3474,24 @@ void InnerLoopVectorizer::fixFixedOrderRecurrence(VPLiveOut *LO,
                                                   VPTransformState &State) {
   // Extract the last vector element in the middle block. This will be the
   // initial value for the recurrence when jumping to the scalar loop.
-  Value *ExtractForScalar = State.get(LO->getOperand(0), UF - 1, true);
+  VPValue *VPExtract = LO->getOperand(0);
+using namespace llvm::VPlanPatternMatch;
+  assert(match(VPExtract, m_VPInstruction<VPInstruction::ExtractFromEnd>(m_VPValue(), m_VPValue())) && "FOR LiveOut expects to use an extract from end.");
+  Value *ResumeScalarFOR = State.get(VPExtract, UF - 1, true);
 
   // Fix the initial value of the original recurrence in the scalar loop.
   Builder.SetInsertPoint(LoopScalarPreHeader, LoopScalarPreHeader->begin());
-  PHINode *Phi = LO->getPhi();
-  auto *Start = Builder.CreatePHI(Phi->getType(), 2, "scalar.recur.init");
-  auto *ScalarInit =
-      LO->getPhi()->getIncomingValueForBlock(LoopScalarPreHeader);
+  PHINode *ScalarHeaderPhi = LO->getPhi();
+  auto *NewScalarHeaderPhi = Builder.CreatePHI(ScalarHeaderPhi->getType(), 2, "scalar.recur.init");
+  auto *InitScalarFOR =
+      ScalarHeaderPhi->getIncomingValueForBlock(LoopScalarPreHeader);
   for (auto *BB : predecessors(LoopScalarPreHeader)) {
-    auto *Incoming = BB == LoopMiddleBlock ? ExtractForScalar : ScalarInit;
-    Start->addIncoming(Incoming, BB);
+    auto *Incoming = BB == LoopMiddleBlock ? ResumeScalarFOR : InitScalarFOR;
+    NewScalarHeaderPhi ->addIncoming(Incoming, BB);
   }
 
-  Phi->setIncomingValueForBlock(LoopScalarPreHeader, Start);
-  Phi->setName("scalar.recur");
+  ScalarHeaderPhi->setIncomingValueForBlock(LoopScalarPreHeader, NewScalarHeaderPhi );
+  ScalarHeaderPhi->setName("scalar.recur");
 }
 
 void InnerLoopVectorizer::sinkScalarOperands(Instruction *PredInst) {

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -3166,7 +3166,7 @@ class VPlan {
   /// definitions are VPValues that hold a pointer to their underlying IR.
   SmallVector<VPValue *, 16> VPLiveInsToFree;
 
-  /// Values used outside the plan.
+  /// Values used outside the plan. It contains live-outs that need fixing. Any live-out that is fixed outside VPlan needs to be removed. The remaining live-outs are fixed via VPLiveOut::fixPhi.
   MapVector<PHINode *, VPLiveOut *> LiveOuts;
 
   /// Mapping from SCEVs to the VPValues representing their expansions.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -886,30 +886,44 @@ bool VPlanTransforms::adjustFixedOrderRecurrences(VPlan &Plan,
     //     br cond, vector.body, middle.block
     //
     //   middle.block:
-    //     x = v2(3)
-    //     br scalar.ph
+    //     s_penultimate = v2(2) = v3(3)
+		//     s_resume = v2(3)
+    //     br cond, scalar.ph, exit.block
     //
     //   scalar.ph:
-    //     s_init = phi [x, middle.block], [a[-1], otherwise]
+    //     s_init' = phi [s_resume, middle.block], [s_init, otherwise]
     //     br scalar.body
     //
+    //   scalar.body:
+    //     i = phi [0, scalar.ph], [i+1, scalar.body]
+    //     s1 = phi [s_init', scalar.ph], [s2, scalar.body]
+    //     s2 = a[i]
+    //     b[i] = s2 - s1
+    //     br cond, scalar.body, exit.block
+    //
+    //   exit.block:
+    //     lo = lcssa.phi [s1, scalar.body], [s.penultimate, middle.block]
+     //
     // After execution completes the vector loop, we extract the next value of
     // the recurrence (x) to use as the initial value in the scalar loop. This
-    // is modeled by ExtractRecurrenceResume.
+    // is modeled by ExtractFromEnd.
     Type *IntTy = Plan.getCanonicalIV()->getScalarType();
-    auto *Result = cast<VPInstruction>(MiddleBuilder.createNaryOp(
+
+    // Extract the penultimate value of the recurrence and update VPLiveOut users of the recurrence splice.
+    auto *Penultimate = cast<VPInstruction>(MiddleBuilder.createNaryOp(
         VPInstruction::ExtractFromEnd,
         {FOR->getBackedgeValue(),
          Plan.getOrAddLiveIn(ConstantInt::get(IntTy, 2))},
         {}, "vector.recur.extract.for.phi"));
     RecurSplice->replaceUsesWithIf(
-        Result, [](VPUser &U, unsigned) { return isa<VPLiveOut>(&U); });
+        Penultimate, [](VPUser &U, unsigned) { return isa<VPLiveOut>(&U); });
+
+    // Extract the resume value and create a new VPLiveOut for it.
     auto *Resume = MiddleBuilder.createNaryOp(
         VPInstruction::ExtractFromEnd,
         {FOR->getBackedgeValue(),
          Plan.getOrAddLiveIn(ConstantInt::get(IntTy, 1))},
         {}, "vector.recur.extract");
-    // Introduce VPUsers modeling the exit values.
     Plan.addLiveOut(cast<PHINode>(FOR->getUnderlyingInstr()), Resume);
   }
   return true;