2 files changed, 20 insertions, 5 deletions

diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
index b5e6418ad32..444bc16e0a1 100644
--- a/llvm/lib/Transforms/Utils/LoopUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -553,13 +553,23 @@ bool RecurrenceDescriptor::isFirstOrderRecurrence(PHINode *Phi, Loop *TheLoop,
   if (!Previous || !TheLoop->contains(Previous) || isa<PHINode>(Previous))
     return false;
 
-  // Ensure every user of the phi node is dominated by the previous value. The
-  // dominance requirement ensures the loop vectorizer will not need to
-  // vectorize the initial value prior to the first iteration of the loop.
   for (User *U : Phi->users())
-    if (auto *I = dyn_cast<Instruction>(U))
+    if (auto *I = dyn_cast<Instruction>(U)) {
+      // Ensure every user of the phi node is dominated by the previous value.
+      // The dominance requirement ensures the loop vectorizer will not need to
+      // vectorize the initial value prior to the first iteration of the loop.
       if (!DT->dominates(Previous, I))
         return false;
+      // When the phi node has users outside the loop, the current logic for
+      // fixFirstOrderRecurrences may generate incorrect code. Specifically, we
+      // extract the last element from the vectorized phi, which would be the
+      // update to the phi before exiting the loop. However, what we want is the
+      // previous phi value before the update (i.e. the second last update
+      // before end of the vectorized loop).
+      // See added test cases in first-order-recurrence.ll
+      if (!TheLoop->contains(I))
+        return false;
+    }
 
   return true;
 }
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 814baca35f6..fff1e29ecbb 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4069,7 +4069,12 @@ void InnerLoopVectorizer::fixFirstOrderRecurrence(PHINode *Phi) {
   VecPhi->addIncoming(Incoming, LI->getLoopFor(LoopVectorBody)->getLoopLatch());
 
   // Extract the last vector element in the middle block. This will be the
-  // initial value for the recurrence when jumping to the scalar loop.
+  // initial value for the recurrence when jumping to the scalar loop. 
+  // FIXME: Note that the last vector element need not always be the correct one:
+  // consider a loop  where we have phi uses outside the loop - we need the
+  // second last iteration value and not the last one). For now, we avoid
+  // considering such cases as firstOrderRecurrences (see
+  // isFirstOrderRecurrence).
   auto *Extract = Incoming;
   if (VF > 1) {
     Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());