[LV][NFC] Improve complexity of fixing users of recurrences

The original loop has O(MxN) since `is_contained` iterates over
all incoming values. This change makes it so only the phis
which use the value as an incoming value are iterated over so
it is now O(M).

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

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index d5b22f9..51ea1a7 100644 (file)
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -3892,12 +3892,16 @@ void InnerLoopVectorizer::fixFixedOrderRecurrence(
   // 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))
-    for (PHINode &LCSSAPhi : LoopExitBlock->phis())
-      if (llvm::is_contained(LCSSAPhi.incoming_values(), Phi)) {
-        LCSSAPhi.addIncoming(ExtractForPhiUsedOutsideLoop, LoopMiddleBlock);
-        State.Plan->removeLiveOut(&LCSSAPhi);
-      }
+  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);
+    }
+  }
 }
 
 void InnerLoopVectorizer::fixReduction(VPReductionPHIRecipe *PhiR,