1 files changed, 37 insertions, 15 deletions

diff --git a/llvm/lib/Analysis/IVDescriptors.cpp b/llvm/lib/Analysis/IVDescriptors.cpp
index 6fb600114bc61..ac81cba836f89 100644
--- a/llvm/lib/Analysis/IVDescriptors.cpp
+++ b/llvm/lib/Analysis/IVDescriptors.cpp
@@ -22,7 +22,6 @@
 #include "llvm/Analysis/MustExecute.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
-#include "llvm/Analysis/ScalarEvolutionExpander.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -699,25 +698,48 @@ bool RecurrenceDescriptor::isFirstOrderRecurrence(
   // 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.
-  // TODO: Consider extending this sinking to handle other kinds of instructions
-  // and expressions, beyond sinking a single cast past Previous.
+  // TODO: Consider extending this sinking to handle memory instructions and
+  // phis with multiple users.
+
+  // Returns true, if all users of I are dominated by DominatedBy.
+  auto allUsesDominatedBy = [DT](Instruction *I, Instruction *DominatedBy) {
+    return all_of(I->uses(), [DT, DominatedBy](Use &U) {
+      return DT->dominates(DominatedBy, U);
+    });
+  };
+
   if (Phi->hasOneUse()) {
-    auto *I = Phi->user_back();
-    if (I->isCast() && (I->getParent() == Phi->getParent()) && I->hasOneUse() &&
-        DT->dominates(Previous, I->user_back())) {
-      if (!DT->dominates(Previous, I)) // Otherwise we're good w/o sinking.
-        SinkAfter[I] = Previous;
+    Instruction *I = Phi->user_back();
+
+    // If the user of the PHI is also the incoming value, we potentially have a
+    // reduction and which cannot be handled by sinking.
+    if (Previous == I)
+      return false;
+
+    // We cannot sink terminator instructions.
+    if (I->getParent()->getTerminator() == I)
+      return false;
+
+    // Do not try to sink an instruction multiple times (if multiple operands
+    // are first order recurrences).
+    // TODO: We can support this case, by sinking the instruction after the
+    // 'deepest' previous instruction.
+    if (SinkAfter.find(I) != SinkAfter.end())
+      return false;
+
+    if (DT->dominates(Previous, I)) // We already are good w/o sinking.
       return true;
-    }
-  }
 
-  for (User *U : Phi->users())
-    if (auto *I = dyn_cast<Instruction>(U)) {
-      if (!DT->dominates(Previous, I))
-        return false;
+    // We can sink any instruction without side effects, as long as all users
+    // are dominated by the instruction we are sinking after.
+    if (I->getParent() == Phi->getParent() && !I->mayHaveSideEffects() &&
+        allUsesDominatedBy(I, Previous)) {
+      SinkAfter[I] = Previous;
+      return true;
     }
+  }
 
-  return true;
+  return allUsesDominatedBy(Phi, Previous);
 }
 
 /// This function returns the identity element (or neutral element) for