[llvm] 175461a - [NFC][LoopVectorize] Make replaceVPBBWithIRVPBB more efficient (#111514)

[via llvm-commits]

Author: David Sherwood
Date: 2024-10-15T14:11:55+01:00
New Revision: 175461a22a72d3407aac5482f1d602dd3e6cb753

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

LOG: [NFC][LoopVectorize] Make replaceVPBBWithIRVPBB more efficient (#111514)

In replaceVPBBWithIRVPBB we spend time erasing and appending
predecessors and successors from a list, when all we really have to do
is replace the old with the new. Not only is this more efficient, but it
also preserves the ordering of successors and predecessors. This is
something which may become important for vectorising early exit loops
(see PR #88385), since a VPIRInstruction is the wrapper for a live-out
phi with extra operands that map to the incoming block according to the
block's predecessor.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VPlan.cpp
    llvm/lib/Transforms/Vectorize/VPlan.h

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 5e3a6388094940..7ded51d9e3abd1 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -1004,13 +1004,9 @@ static void replaceVPBBWithIRVPBB(VPBasicBlock *VPBB, BasicBlock *IRBB) {
     assert(!R.isPhi() && "Tried to move phi recipe to end of block");
     R.moveBefore(*IRVPBB, IRVPBB->end());
   }
-  VPBlockBase *PredVPBB = VPBB->getSinglePredecessor();
-  VPBlockUtils::disconnectBlocks(PredVPBB, VPBB);
-  VPBlockUtils::connectBlocks(PredVPBB, IRVPBB);
-  for (auto *Succ : to_vector(VPBB->getSuccessors())) {
-    VPBlockUtils::connectBlocks(IRVPBB, Succ);
-    VPBlockUtils::disconnectBlocks(VPBB, Succ);
-  }
+
+  VPBlockUtils::reassociateBlocks(VPBB, IRVPBB);
+
   delete VPBB;
 }
 

diff  --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 6a61ef63c2a054..ec91d08db5ba00 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -440,6 +440,26 @@ class VPBlockBase {
     Successors.erase(Pos);
   }
 
+  /// This function replaces one predecessor with another, useful when
+  /// trying to replace an old block in the CFG with a new one.
+  void replacePredecessor(VPBlockBase *Old, VPBlockBase *New) {
+    auto I = find(Predecessors, Old);
+    assert(I != Predecessors.end());
+    assert(Old->getParent() == New->getParent() &&
+           "replaced predecessor must have the same parent");
+    *I = New;
+  }
+
+  /// This function replaces one successor with another, useful when
+  /// trying to replace an old block in the CFG with a new one.
+  void replaceSuccessor(VPBlockBase *Old, VPBlockBase *New) {
+    auto I = find(Successors, Old);
+    assert(I != Successors.end());
+    assert(Old->getParent() == New->getParent() &&
+           "replaced successor must have the same parent");
+    *I = New;
+  }
+
 protected:
   VPBlockBase(const unsigned char SC, const std::string &N)
       : SubclassID(SC), Name(N) {}
@@ -3947,6 +3967,19 @@ class VPBlockUtils {
     To->removePredecessor(From);
   }
 
+  /// Reassociate all the blocks connected to \p Old so that they now point to
+  /// \p New.
+  static void reassociateBlocks(VPBlockBase *Old, VPBlockBase *New) {
+    for (auto *Pred : to_vector(Old->getPredecessors()))
+      Pred->replaceSuccessor(Old, New);
+    for (auto *Succ : to_vector(Old->getSuccessors()))
+      Succ->replacePredecessor(Old, New);
+    New->setPredecessors(Old->getPredecessors());
+    New->setSuccessors(Old->getSuccessors());
+    Old->clearPredecessors();
+    Old->clearSuccessors();
+  }
+
   /// Return an iterator range over \p Range which only includes \p BlockTy
   /// blocks. The accesses are casted to \p BlockTy.
   template <typename BlockTy, typename T>