[llvm] f5fe7ab - [VPlan] Account for removed users in replaceAllUsesWith.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2020-10-04T18:18:58+01:00
New Revision: f5fe7abe8a8c8150b7b305bae963b429f15ea217

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

LOG: [VPlan] Account for removed users in replaceAllUsesWith.

Make sure we do not iterate using an invalid iterator.

Another small fix/step towards traversing the def-use chains in VPlan.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VPlan.cpp
    llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 417983241023..95d5cfdcafba 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -884,10 +884,18 @@ void VPWidenCanonicalIVRecipe::print(raw_ostream &O, const Twine &Indent,
 template void DomTreeBuilder::Calculate<VPDominatorTree>(VPDominatorTree &DT);
 
 void VPValue::replaceAllUsesWith(VPValue *New) {
-  for (VPUser *User : users())
+  for (unsigned J = 0; J < getNumUsers();) {
+    VPUser *User = Users[J];
+    unsigned NumUsers = getNumUsers();
     for (unsigned I = 0, E = User->getNumOperands(); I < E; ++I)
       if (User->getOperand(I) == this)
         User->setOperand(I, New);
+    // If a user got removed after updating the current user, the next user to
+    // update will be moved to the current position, so we only need to
+    // increment the index if the number of users did not change.
+    if (NumUsers == getNumUsers())
+      J++;
+  }
 }
 
 void VPValue::printAsOperand(raw_ostream &OS, VPSlotTracker &Tracker) const {

diff  --git a/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
index a64f9e374ebf..926473af7d46 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
@@ -167,7 +167,13 @@ TEST(VPInstructionTest, replaceAllUsesWith) {
   EXPECT_EQ(0u, VPV2->getNumUsers());
   EXPECT_EQ(0u, VPV3->getNumUsers());
 
+  VPInstruction *I2 = new VPInstruction(0, {VPV1, VPV2});
+  EXPECT_EQ(3u, VPV1->getNumUsers());
+  VPV1->replaceAllUsesWith(VPV3);
+  EXPECT_EQ(3u, VPV3->getNumUsers());
+
   delete I1;
+  delete I2;
   delete VPV1;
   delete VPV2;
   delete VPV3;