[PATCH] D147944: [LV][NFC] Improve complexity of fixing users of reductions

[Michael Maitland via Phabricator via llvm-commits]

michaelmaitland updated this revision to Diff 512180.
michaelmaitland marked 2 inline comments as done.
michaelmaitland added a comment.

Use auto and dyn_cast


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D147944/new/

https://reviews.llvm.org/D147944

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


Index: llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
===================================================================
--- llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4100,12 +4100,18 @@
   // We know that the loop is in LCSSA form. We need to update the PHI nodes
   // in the exit blocks.  See comment on analogous loop in
   // fixFixedOrderRecurrence for a more complete explaination of the logic.
-  if (!Cost->requiresScalarEpilogue(VF))
-    for (PHINode &LCSSAPhi : LoopExitBlock->phis())
-      if (llvm::is_contained(LCSSAPhi.incoming_values(), LoopExitInst)) {
-        LCSSAPhi.addIncoming(ReducedPartRdx, LoopMiddleBlock);
-        State.Plan->removeLiveOut(&LCSSAPhi);
-      }
+  if (!Cost->requiresScalarEpilogue(VF)) {
+    SmallPtrSet<PHINode *, 2> ToFix;
+    auto *Phi = cast<PHINode>(PhiR->getUnderlyingValue());
+    for (User *U : Phi->users())
+      if (auto *UPhi = dyn_cast<PHINode>(U);
+          UPhi && UPhi->getParent() == LoopExitBlock)
+        ToFix.insert(UPhi);
+    for (PHINode *LCSSAPhi : ToFix) {
+      LCSSAPhi->addIncoming(ReducedPartRdx, LoopMiddleBlock);
+      State.Plan->removeLiveOut(LCSSAPhi);
+    }
+  }
 
   // Fix the scalar loop reduction variable with the incoming reduction sum
   // from the vector body and from the backedge value.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D147944.512180.patch
Type: text/x-patch
Size: 1355 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230410/c600145f/attachment.bin>