[llvm] e86ed9b - [LV][NFC] Improve complexity of fixing users of recurrences

[Michael Maitland via llvm-commits]

Author: Michael Maitland
Date: 2023-04-06T16:15:51-07:00
New Revision: e86ed9bf2a61409553ba51bf218bd97d2b9e9132

URL: https://github.com/llvm/llvm-project/commit/e86ed9bf2a61409553ba51bf218bd97d2b9e9132
DIFF: https://github.com/llvm/llvm-project/commit/e86ed9bf2a61409553ba51bf218bd97d2b9e9132.diff

LOG: [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

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index d5b22f9484a0b..51ea1a70caa57 100644
--- 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,